from flask import Flask, render_template, request, jsonify
import subprocess
import json
import os
import time
import copy
import random
import string
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"

BASE_REQUESTS = [
    {"id": 1, "type": "malicious", "cmd": "curl -A '1' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"query={{{DYNAMIC}{DYNAMIC}{DYNAMIC}}}\" \"http://localhost:8080/search\""},
    {"id": 2, "type": "malicious", "cmd": "curl -A '2' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"path={{{DYNAMIC}}}/test\" \"http://localhost:8080/\""},
    {"id": 3, "type": "malicious", "cmd": "curl -A '3' -H \"X-Forwarded-For: #{{DYNAMIC}{DYNAMIC}{DYNAMIC}}\" 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\":\"[%{DYNAMIC}-{DYNAMIC}+{DYNAMIC}%]\",\"password\":\"x\"}'"},
    {"id": 5, "type": "malicious", "cmd": "curl -A '5' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"username=<%'{DYNAMIC}*{DYNAMIC}%>'&password=anything\" \"http://localhost:8080/login\""},
    {"id": 6, "type": "malicious", "cmd": "curl -A '6' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"id={%{DYNAMIC}%}\" \"http://localhost:8080/item\""},
    {"id": 7, "type": "malicious", "cmd": "curl -A '7' -H \"Cookie: session_id={{config.SECRET_KEY}}\" http://localhost:8080/dashboard"},
    {"id": 8, "type": "malicious", "cmd": "curl -A '8' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"filter={{{DYNAMIC}}}\" \"http://localhost:8080/data\""},
    {"id": 9, "type": "malicious", "cmd": "curl -A '9' -H \"X-User-Key: @{{DYNAMIC}{DYNAMIC}{DYNAMIC}}\" \"http://localhost:8080/search\""},
    {"id": 10, "type": "malicious", "cmd": "curl -A '10' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"id={{[].__class__.__base__.__subclasses__()[80].__init__.__globals__['__builtins__'].eval('__import__(\\\"os\\\").system(\\\"ls\\\")')}}\" \"http://localhost:8080/profile\""},
    {"id": 11, "type": "malicious", "cmd": "curl -A '11' -X POST http://localhost:8080/aspnet/api/contact -H 'Content-Type: application/json' -d '{\"message\":\"<%= {DYNAMIC} %>\"}'"},
    {"id": 12, "type": "malicious", "cmd": "curl -A '12' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"id=$[{DYNAMIC}{DYNAMIC}{DYNAMIC}]\" \"http://localhost:8080/profile\""},

    {"id": 16, "type": "normal", "cmd": "curl -A '16' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"query[{DYNAMIC_NORM}]=O'Reilly&query[{DYNAMIC_NORM}]=O'Brien\" \"http://localhost:8080/search\""},
    {"id": 17, "type": "normal", "cmd": "curl -A '17' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"username[{DYNAMIC_NORM}]=test&password[{DYNAMIC_NORM}]=){}123%^\" \"http://localhost:8080/login\""},
    {"id": 18, "type": "normal", "cmd": "curl -A '18' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"id={DYNAMIC_NORM} OR {DYNAMIC_NORM}\" \"http://localhost:8080/user\""},
    {"id": 19, "type": "normal", "cmd": "curl -A '19' -H \"Cookie: jsessid=ciqw891; token=1789971|{791}|{}|1898\" http://localhost:8080/api/user"},
    {"id": 20, "type": "normal", "cmd": "curl -A '20' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"where={DYNAMIC_NORM} FROM users\" \"http://localhost:8080/filter\""},
    {"id": 21, "type": "normal", "cmd": "curl -A '21' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"id={{DYNAMIC_NORM}}\" \"http://localhost:8080/file\""},
    {"id": 22, "type": "normal", "cmd": "curl -A '22' -H \"X-User-Name: {DYNAMIC_NORM}%_C{_%8>\" \"http://localhost:8080/\""},
    {"id": 23, "type": "normal", "cmd": "curl -A '23' -X POST -H \"Content-Type: application/x-www-form-urlencoded\" --data-binary \"id={DYNAMIC_NORM}{DYNAMIC_NORM}{DYNAMIC_NORM}\" \"http://localhost:8080/upload\""},
    {"id": 24, "type": "normal", "cmd": "curl -A '24' -X POST http://localhost:8080/php/contact.php -d \"name=John&email=john@example.com&message=Hello, my invoice is {DYNAMIC_NORM} and I need {DYNAMIC_NORM} refund\""},
]

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

suricata_process = None
test_server_process = None

def generate_dynamic_value():
    choice = random.choice(['number', 'string', 'symbol', 'word'])
    if choice == 'number':
        return str(random.randint(1, 100))
    elif choice == 'string':
        length = random.randint(1, 3)
        return ''.join(random.choices(string.ascii_letters + string.digits, k=length))
    elif choice == 'word':
        return random.choice(['*', '-', '+', '/'])
    else:
        return random.choice(['env', 'config', 'settings.config()', 'environ'])

def generate_dynamic_value_norm():
    choice = random.choice(['number', 'sum', 'word'])
    if choice == 'number':
        return str(random.randint(1, 100))
    elif choice == 'sum':
        return f"${random.randint(1, 100)}"
    else:
        return random.choice(['OR', 'AND'])

def inject_dynamic_into_requests(requests_list):
    generated = []

    for req in requests_list:
        new_req = copy.deepcopy(req)
        dynamic_value = generate_dynamic_value()
        dynamic_value_norm = generate_dynamic_value_norm()
        if '{DYNAMIC}' in new_req['cmd'] or '{DYNAMIC_NORM}' in new_req['cmd']:
            new_req['cmd'] = new_req['cmd'].replace('{DYNAMIC}', dynamic_value)
            new_req['cmd'] = new_req['cmd'].replace('{DYNAMIC_NORM}', dynamic_value_norm)
        if isinstance(new_req['id'], str) and ('{DYNAMIC}' in new_req['id'] or '{DYNAMIC_NORM}' in new_req['id']):
            try:
                new_req['id'] = int(new_req['id'].replace('{DYNAMIC}', dynamic_value))
                new_req['id'] = int(new_req['id'].replace('{DYNAMIC_NORM}', dynamic_value_norm))
            except:
                new_req['id'] = new_req['id'].replace('{DYNAMIC}', dynamic_value)
                new_req['id'] = new_req['id'].replace('{DYNAMIC_NORM}', dynamic_value_norm)
        new_req['dynamic_used'] = dynamic_value
        generated.append(new_req)

    return generated

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)

    global REQUESTS
    REQUESTS = inject_dynamic_into_requests(BASE_REQUESTS)

    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)