import os import json import subprocess import pandas as pd import xgboost as xgb from flask import Flask, request, render_template, redirect, url_for, send_from_directory, flash, jsonify from werkzeug.utils import secure_filename from collections import Counter # --- Configuration --- # Use absolute paths for robustness, assuming app.py is in the project root PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__)) UPLOAD_FOLDER = os.path.join(PROJECT_ROOT, 'uploads') REPORTS_FOLDER = os.path.join(PROJECT_ROOT, 'data', 'reports') # Import your feature extractor # This assumes feature_extractor.py is in a 'scripts' subfolder. # If your structure is different, you might need to adjust the import. try: from scripts.feature_extractor import extract_features_from_json except ImportError: # A fallback if the script is not found, to prevent the app from crashing on start. def extract_features_from_json(path): print(f"WARNING: Could not import 'extract_features_from_json'. Using a dummy function.") # Return a dictionary with expected feature names and default values. # This allows the prediction part to run without error, though results will be meaningless. return {f'feature_{i}': 0 for i in range(100)} # Paths to scripts and models ANALYZER_SCRIPT = os.path.join(PROJECT_ROOT, 'analyzer', 'forensic_apk_analyzer.py') MODEL_PATH = os.path.join(PROJECT_ROOT, 'models', 'xgboost_model.json') ALLOWED_EXTENSIONS = {'apk'} app = Flask(__name__) app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER app.config['REPORTS_FOLDER'] = REPORTS_FOLDER app.secret_key = 'a-very-secret-key-for-production' # --- Pre-load ML Model --- try: model = xgb.XGBClassifier() model.load_model(MODEL_PATH) except Exception as e: model = None print(f"CRITICAL: Could not load XGBoost model from {MODEL_PATH}. Prediction will fail. Error: {e}") # --- Helper Functions --- def allowed_file(filename): """Checks if the uploaded file has an allowed extension.""" return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS def create_folders(): """Creates necessary folders for uploads and reports if they don't exist.""" os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) os.makedirs(app.config['REPORTS_FOLDER'], exist_ok=True) # --- Flask Routes --- @app.route('/', methods=['GET']) def index(): """Renders the main upload page.""" return render_template('index.html') @app.route('/analyze', methods=['POST']) def analyze(): """Handles file upload, runs analysis, and performs AI prediction.""" if model is None: return jsonify({'error': 'AI model is not loaded. Check server logs.'}), 500 if 'file' not in request.files: return jsonify({'error': 'No file part in the request.'}), 400 file = request.files['file'] if file.filename == '': return jsonify({'error': 'No file selected.'}), 400 if file and allowed_file(file.filename): filename = secure_filename(file.filename) apk_path = os.path.join(app.config['UPLOAD_FOLDER'], filename) file.save(apk_path) # --- Step 1: Run Forensic Analyzer --- try: # Execute the forensic analyzer script with a timeout. subprocess.run( ['python', ANALYZER_SCRIPT, apk_path, '-o', app.config['REPORTS_FOLDER']], check=True, capture_output=True, text=True, timeout=300 # 5-minute timeout ) report_filename = f"{os.path.splitext(filename)[0]}_forensic_report.json" report_path = os.path.join(app.config['REPORTS_FOLDER'], report_filename) if not os.path.exists(report_path): return jsonify({'error': 'Analysis completed, but the report file was not found.'}), 500 except subprocess.CalledProcessError as e: return jsonify({'error': f"Analysis script failed with an error: {e.stderr}"}), 500 except subprocess.TimeoutExpired: return jsonify({'error': 'Analysis timed out. The APK might be too large or complex.'}), 408 # --- Step 2: Extract Features & Predict --- try: features_dict = extract_features_from_json(report_path) # Convert features to a DataFrame for the model. X_new = pd.DataFrame([features_dict]) prediction_val = model.predict(X_new)[0] prediction_label = 'High-Risk / Malware' if prediction_val == 1 else 'Benign / Safe' except Exception as e: return jsonify({'error': f"Feature extraction or prediction failed: {str(e)}"}), 500 # --- Success --- # Return the report URL and the prediction label. return jsonify({ 'report_url': url_for('report', filename=report_filename), 'prediction': prediction_label }) return jsonify({'error': 'Invalid file type. Only .apk files are permitted.'}), 400 @app.route('/report/') def report(filename): """Displays the detailed analysis report.""" report_path = os.path.join(app.config['REPORTS_FOLDER'], filename) try: with open(report_path, 'r', encoding='utf-8') as f: report_data = json.load(f) except (FileNotFoundError, json.JSONDecodeError): flash(f"Could not read or parse the report file: {filename}") return redirect(url_for('index')) # Process geolocation data for the map geolocated_ips = report_data.get('code_analysis', {}).get('geolocated_ips', {}) locations = [] location_counts = Counter() if geolocated_ips: for ip, data in geolocated_ips.items(): if data and 'loc' in data and 'error' not in data: try: lat, lon = map(float, data['loc'].split(',')) city = data.get('city', 'Unknown City') country = data.get('country', 'Unknown Country') location_key = f"{city}, {country}" locations.append({'lat': lat, 'lon': lon, 'ip': ip, 'city': city, 'country': country}) location_counts[location_key] += 1 except (ValueError, AttributeError): # Skip if location data is malformed. continue # Identify the top 3 most common locations. top_locations_set = {loc for loc, count in location_counts.most_common(3)} for loc in locations: loc['top_location'] = f"{loc['city']}, {loc['country']}" in top_locations_set return render_template( 'report.html', report=report_data, filename=filename, locations=json.dumps(locations) # Pass locations as a JSON string. ) @app.route('/download/') def download_report(filename): """Allows the user to download the full JSON report.""" return send_from_directory( app.config['REPORTS_FOLDER'], filename, as_attachment=True ) if __name__ == '__main__': create_folders() app.run(debug=True)