web.py

from flask import Flask, request, render_template
import os
import predict

app = Flask(__name__)

# Set the upload folder and allowed file extensions
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'exe', 'hta'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Function to check if a file has an allowed extension


def predictMalware(filename):
    # Path to the software file you want to test
    # software_path = filename
    software_path = "D:\Scan\TeamViewer_Setup_x64.exe"
    print(f"file name: {filename}\nsoftware path: {software_path}")
    # Extract features from the software file
    extracted_features = predict.extract_features(software_path)
    print(extracted_features)

    if extracted_features is not None:
        # Preprocess the 'BaseOfData' feature to handle 'inf' values
        extracted_features = [predict.sys.float_info.max if x ==
                              float('inf') else x for x in extracted_features]

        # Make a prediction using the loaded model
        prediction = predict.loaded_classifier.predict([extracted_features])

        # Interpret the prediction
        message = ""
        if prediction[0] == 1:
            message = "The software is predicted to be legitimate."
            print("The software is predicted to be legitimate.")
        else:
            message = "The software is predicted to be malicious."
            print("The software is predicted to be malicious.")
    else:
        # Feature extraction failed
        print("Feature extraction failed.")

    return message

# Route to handle user feedback


# @app.route('/setFeedback', methods=['POST'])
# def feedback():
#     #


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        # Check if a file was uploaded
        if 'file' not in request.files:
            return render_template('upload.html', message='No file part')

        file = request.files['file']

        # Check if the file is empty
        if file.filename == '':
            return render_template('upload.html', message='No selected file')

        # Check if the file has an allowed extension
        if file and allowed_file(file.filename):
            # Save the uploaded file to the UPLOAD_FOLDER
            # ts = str(time.time())
            filename = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
            file.save(filename)
            predictMalware(filename)
            return render_template('upload.html', message='File uploaded successfully.', Prediction=f' Prediction: {predictMalware(filename)}')

    return render_template('upload.html')


@app.route('/confirm', methods=['GET', 'POST'])
def confirm():
    if request.method == 'GET':
        return render_template('confirmation.html')
    # Handle the confirmation result here (you can save it to a database or take appropriate action)
    confirmation = request.form.get('confirmation')
    if confirmation == 'correct':
        # Process as correct prediction
        pass
    elif confirmation == 'wrong':
        # Process as incorrect prediction
        pass

    return "Confirmation received."


if __name__ == '__main__':
    # Ensure the "uploads" directory exists
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)

    # run this app run glo
    app.run(host="0.0.0.0")