Add files via upload

Author: 0xfke

app.py

@@ -0,0 +1,56 @@
+from flask import Flask, request, render_template
+import joblib
+import os
+from feature_extraction import extract_features
+
+app = Flask(__name__)
+
+# Load the trained model
+model = joblib.load('ML_model/malwareclassifier-V2.pkl')
+
+UPLOAD_FOLDER = 'uploads'
+if not os.path.exists(UPLOAD_FOLDER):
+    os.makedirs(UPLOAD_FOLDER)
+
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+ALLOWED_EXTENSIONS = {'dll', 'exe'}
+
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+@app.route('/')
+def index():
+    return render_template('index.html')
+
+@app.route('/analyze', methods=['POST'])
+def analyze():
+    # Check if a file is uploaded
+    if 'file' in request.files:
+        file = request.files['file']
+        
+        if file.filename == '' or not allowed_file(file.filename):
+            return render_template('index.html', error="Unsupported file type.")
+        
+        # Construct the full file path
+        file_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
+
+        # Save the file
+        file.save(file_path)
+
+        # Use the model for prediction if the file is `.exe` or `.dll`
+        if allowed_file(file.filename):
+            features = extract_features(file_path)  # Your feature extraction function
+            prediction = model.predict(features)     # Predict using your model
+            result = {
+                "type": "file",
+                "prediction": "Malware" if prediction[0] == 1 else "Safe",
+                "file_name": file.filename
+            }
+
+        return render_template('result.html', result=result)
+
+    return render_template('index.html', error="No file uploaded.")
+
+if __name__ == '__main__':
+    app.run(port=5001, debug=True)
+

feature_extraction.py

@@ -0,0 +1,58 @@
+import pefile
+import pandas as pd
+import math
+
+# Function to calculate entropy of a section
+def calculate_entropy(data):
+    if not data:
+        return 0
+    entropy = 0
+    for x in range(256):
+        p_x = float(data.count(bytes([x]))) / len(data)
+        if p_x > 0:
+            entropy += - p_x * math.log(p_x, 2)
+    return entropy
+
+def extract_features(file_path):
+    pe = pefile.PE(file_path)
+
+    # Extract the specified 23 features in the given order
+    features = {
+        'MajorLinkerVersion': pe.OPTIONAL_HEADER.MajorLinkerVersion,
+        'MinorOperatingSystemVersion': pe.OPTIONAL_HEADER.MinorOperatingSystemVersion,
+        'MajorSubsystemVersion': pe.OPTIONAL_HEADER.MajorSubsystemVersion,
+        'SizeOfStackReserve': pe.OPTIONAL_HEADER.SizeOfStackReserve,
+        'TimeDateStamp': pe.FILE_HEADER.TimeDateStamp,
+        'MajorOperatingSystemVersion': pe.OPTIONAL_HEADER.MajorOperatingSystemVersion,
+        'Characteristics': pe.FILE_HEADER.Characteristics,
+        'ImageBase': pe.OPTIONAL_HEADER.ImageBase,
+        'Subsystem': pe.OPTIONAL_HEADER.Subsystem,
+        'MinorImageVersion': pe.OPTIONAL_HEADER.MinorImageVersion,
+        'MinorSubsystemVersion': pe.OPTIONAL_HEADER.MinorSubsystemVersion,
+        'SizeOfInitializedData': pe.OPTIONAL_HEADER.SizeOfInitializedData,
+        'DllCharacteristics': pe.OPTIONAL_HEADER.DllCharacteristics,
+        'DirectoryEntryExport': 1 if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') else 0,
+        'ImageDirectoryEntryExport': pe.OPTIONAL_HEADER.DATA_DIRECTORY[0].Size if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') else 0,
+        'CheckSum': pe.OPTIONAL_HEADER.CheckSum,
+        'DirectoryEntryImportSize': pe.OPTIONAL_HEADER.DATA_DIRECTORY[1].Size if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT') else 0,
+        'SectionMaxChar': len(pe.sections),  # Example calculation for demonstration
+        'MajorImageVersion': pe.OPTIONAL_HEADER.MajorImageVersion,
+        'AddressOfEntryPoint': pe.OPTIONAL_HEADER.AddressOfEntryPoint,
+        'SectionMinEntropy': None,  # Placeholder, will be calculated
+        'SizeOfHeaders': pe.OPTIONAL_HEADER.SizeOfHeaders,
+        'SectionMinVirtualsize': None  # Placeholder, will be calculated
+    }
+
+    # Calculate SectionMinEntropy
+    entropies = []
+    for section in pe.sections:
+        entropy = calculate_entropy(section.get_data())
+        entropies.append(entropy)
+
+    if entropies:
+        features['SectionMinEntropy'] = min(entropies)
+
+    # Calculate SectionMinVirtualsize (example calculation)
+    features['SectionMinVirtualsize'] = min(section.Misc_VirtualSize for section in pe.sections)
+
+    return pd.DataFrame([features])