from flask import Flask, render_template, request, jsonify
import subprocess
import json
import os
import time
import re

app = Flask(__name__, static_folder='static')
app.config['SESSION_COOKIE_NAME'] = "task_session"

SURICATA_CONFIG = "/etc/suricata/suricata.yaml"
RULES_FILE = "/etc/suricata/rules/custom.rules"
EVE_LOG = "/var/log/suricata/eve.json"

REQUESTS = [
    {"id": 1, "type": "malicious", "cmd": "curl -A '1' \"http://localhost:8080/search?query=' UNION SELECT username,password FROM users--\""},
    {"id": 2, "type": "malicious", "cmd": "curl -A '2' \"http://localhost:8080/product?id=1' AND SLEEP(5)--\""},
    {"id": 3, "type": "malicious", "cmd": "curl -A '3' -H \"X-Forwarded-For: 127.0.0.1' OR '1'='1'\" http://localhost:8080/admin"},
    {"id": 4, "type": "malicious", "cmd": "curl -A '4' -X POST http://localhost:8080/api/login -H 'Content-Type: application/json' -d '{\"username\":\"admin\" OR 1=1--\",\"password\":\"x\"}'"},
    {"id": 5, "type": "malicious", "cmd": "curl -A '5' -X POST http://localhost:8080/login -d \"username=admin'--&password=anything\""},
    {"id": 6, "type": "malicious", "cmd": "curl -A '6' \"http://localhost:8080/item?id=1' AND SUBSTRING((SELECT password FROM users LIMIT 1),1,1)='a'--\""},
    {"id": 7, "type": "malicious", "cmd": "curl -A '7' -H \"Cookie: session_id=' OR (SELECT COUNT(*) FROM users)>0--\" http://localhost:8080/dashboard"},
    {"id": 8, "type": "malicious", "cmd": "curl -A '8' \"http://localhost:8080/data?filter=' OR 1=1--\""},
    {"id": 9, "type": "malicious", "cmd": "curl -A '9' \"http://localhost:8080/search?q=1';SELECT * FROM users;--\""},
    {"id": 10, "type": "malicious", "cmd": "curl -A '10' \"http://localhost:8080/profile?id=1' AND EXTRACTVALUE(1,CONCAT(0x7e,(SELECT version()),0x7e))--\""},
    {"id": 16, "type": "normal", "cmd": "curl -A '16' \"http://localhost:8080/search?query=O'Reilly\""},
    {"id": 17, "type": "normal", "cmd": "curl -A '17' -X POST http://localhost:8080/login -d \"username=O'Connor&password=securePass123\""},
    {"id": 18, "type": "normal", "cmd": "curl -A '18' \"http://localhost:8080/user?id=100 OR 200\""},
    {"id": 19, "type": "normal", "cmd": "curl -A '19' -H \"X-User-Data: test'-data\" http://localhost:8080/api/user"},
    {"id": 20, "type": "normal", "cmd": "curl -A '20' \"http://localhost:8080/filter?where=name FROM users\""},
    {"id": 21, "type": "normal", "cmd": "curl -A '21' \"http://localhost:8080/upload?id='1'\""},
]

SECRET_FLAG = os.environ.get('FLAG', 'SURICATA_PERFECT_DETECTION_2024')

suricata_process = None
test_server_process = None

def get_current_rules():
    try:
        with open(RULES_FILE, 'r') as f:
            return f.read()
    except:
        return ""

def save_rules(rules_text):
    with open(RULES_FILE, 'w') as f:
        f.write(rules_text)

def validate_rules_count(rules_text):
    arrow_count = rules_text.count("->")
    if arrow_count > 3:
        return False, f"Слишком много правил! Максимум: 3"
    return True, ""

def validate_rules_syntax():
    result = subprocess.run(
        ["suricata", "-T", "-c", SURICATA_CONFIG],
        capture_output=True,
        text=True,
        timeout=10
    )
    return result.returncode == 0, result.stdout + result.stderr

def clear_logs():
    if os.path.exists(EVE_LOG):
        os.remove(EVE_LOG)

def start_test_server():
    global test_server_process
    test_server_process = subprocess.Popen(
        ["python3", "-m", "http.server", "8080", "--bind", "127.0.0.1"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )
    time.sleep(1)
    return True

def stop_test_server():
    global test_server_process
    if test_server_process:
        test_server_process.terminate()
        test_server_process.wait()

def start_suricata():
    global suricata_process
    stop_suricata()

    subprocess.run(["ip", "link", "set", "dev", "lo", "mtu", "1500"], check=False)

    suricata_process = subprocess.Popen(
        ["suricata", "-c", SURICATA_CONFIG, "-i", "lo"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )
    time.sleep(2)
    return True

def stop_suricata():
    global suricata_process
    if suricata_process:
        suricata_process.terminate()
        suricata_process.wait()

def extract_id_from_alert(alert):
    try:
        if 'http' in alert and 'http_user_agent' in alert['http']:
            user_agent = alert['http']['http_user_agent']
            if user_agent.strip().isdigit():
                return int(user_agent.strip())
            match = re.search(r'(\d+)', user_agent)
            if match:
                return int(match.group(1))
    except:
        pass
    return None

def run_tests(rules_text):
    save_rules(rules_text)

    clear_logs()

    count_valid, count_error = validate_rules_count(rules_text)
    if not count_valid:
        return {"validation_error": count_error}

    valid, validation_output = validate_rules_syntax()
    if not valid:
        return {"validation_error": validation_output}

    start_test_server()
    start_suricata()
    time.sleep(2)

    for req in REQUESTS:
        try:
            subprocess.run(req["cmd"], shell=True, capture_output=True, text=True, timeout=2)
            time.sleep(0.7)
        except:
            pass

    time.sleep(2)

    alert_events = []
    if os.path.exists(EVE_LOG):
        with open(EVE_LOG, 'r') as f:
            for line in f:
                try:
                    event = json.loads(line.strip())
                    if event.get("event_type") == "alert":
                        alert_events.append(event)
                except:
                    continue

    tp_alerts_by_id = {}
    fp_alerts_by_id = {}

    for alert in alert_events:
        request_id = extract_id_from_alert(alert)

        if request_id is not None:
            if request_id <= 15:
                if request_id not in tp_alerts_by_id:
                    alert['detected_request_id'] = request_id
                    alert['detected_type'] = 'TP'
                    tp_alerts_by_id[request_id] = alert
            else:
                if request_id not in fp_alerts_by_id:
                    alert['detected_request_id'] = request_id
                    alert['detected_type'] = 'FP'
                    fp_alerts_by_id[request_id] = alert

    tp_ids = sorted(tp_alerts_by_id.keys())
    fp_ids = sorted(fp_alerts_by_id.keys())

    expected_tp = [r['id'] for r in REQUESTS if r['id'] <= 15]
    expected_fp = [r['id'] for r in REQUESTS if r['id'] > 15]

    perfect_detection = (
        set(tp_ids) == set(expected_tp) and
        len(fp_ids) == 0
    )

    stop_suricata()
    stop_test_server()

    return {
        "tp_ids": tp_ids,
        "fp_ids": fp_ids,
        "tp_count": len(tp_ids),
        "fp_count": len(fp_ids),
        "tp_total": len(expected_tp),
        "fp_total": len(expected_fp),
        "tp_alerts": list(tp_alerts_by_id.values()),
        "fp_alerts": list(fp_alerts_by_id.values()),
        "perfect_detection": perfect_detection,
        "flag": SECRET_FLAG if perfect_detection else None
    }

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/rules', methods=['GET'])
def get_rules():
    return jsonify({"rules": get_current_rules()})

@app.route('/api/test', methods=['POST'])
def test_rules():
    try:
        data = request.json
        rules_text = data.get("rules", "")
        results = run_tests(rules_text)
        return jsonify(results)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)