"""
OmniAnti - Advanced Threat Protection System
Real system scanner backend (Windows)
Run: python app.py
Then open: http://localhost:5000
"""

import os
import sys
import json
import hashlib
import subprocess
import threading
import time
import platform
import re
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
import queue

# ─────────────────────────────────────────────
#  KNOWN MALICIOUS HASHES (SHA-256)
#  A small built-in set — extend this list
#  or load from an external file/API.
# ─────────────────────────────────────────────
KNOWN_BAD_HASHES = {
    # EICAR test virus hash
    "275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f",
    # Add more SHA-256 hashes here
}

# Suspicious file extensions
SUSPICIOUS_EXTENSIONS = {
    ".exe", ".dll", ".bat", ".cmd", ".vbs", ".js", ".ps1",
    ".wsf", ".hta", ".scr", ".pif", ".com", ".jar", ".msi"
}

# Suspicious filename patterns
SUSPICIOUS_PATTERNS = [
    r"crack", r"keygen", r"hack", r"trojan", r"malware",
    r"virus", r"ransom", r"crypt0r", r"worm", r"rootkit",
    r"backdoor", r"keylog", r"stealer", r"inject", r"payload",
    r"dropper", r"exploit", r"~tmp.*\.exe", r"svchost_\d+",
]

# Suspicious startup/registry locations (Windows)
SUSPICIOUS_STARTUP_PATHS = [
    r"C:\Windows\Temp",
    r"C:\Users\Public",
    r"%APPDATA%\Roaming",
    r"%TEMP%",
]

# ─────────────────────────────────────────────
#  SCAN STATE (shared between threads)
# ─────────────────────────────────────────────
scan_state = {
    "running": False,
    "progress": 0,
    "files_scanned": 0,
    "threats": [],
    "logs": [],
    "cmds_done": [],
    "complete": False,
    "scan_type": "quick",
}
log_queue = queue.Queue()


def add_log(msg, level="info"):
    entry = {"time": time.strftime("%H:%M:%S"), "msg": msg, "level": level}
    scan_state["logs"].append(entry)
    log_queue.put(entry)


# ─────────────────────────────────────────────
#  REAL SYSTEM COMMANDS
# ─────────────────────────────────────────────

def run_cmd(cmd, label, timeout=15):
    """Run a real system command and return output."""
    add_log(f"Running: {cmd}", "cmd")
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout
        )
        output = result.stdout.strip() or result.stderr.strip()
        scan_state["cmds_done"].append({"cmd": cmd, "label": label, "output": output[:2000]})
        return output
    except subprocess.TimeoutExpired:
        add_log(f"Command timed out: {cmd}", "warn")
        return ""
    except Exception as e:
        add_log(f"Command error: {e}", "warn")
        return ""


def check_network_connections():
    """Check for suspicious open ports and connections."""
    add_log("Analyzing network connections...", "info")
    out = run_cmd("netstat -ano", "Network Connections")
    suspicious = []
    known_suspicious_ports = {4444, 31337, 1337, 6666, 7777, 8888, 9999, 12345, 54321}

    for line in out.splitlines():
        parts = line.split()
        if len(parts) >= 5:
            try:
                port = int(parts[1].split(":")[-1])
                if port in known_suspicious_ports:
                    suspicious.append(f"Suspicious port open: {parts[1]} (PID: {parts[-1]})")
            except (ValueError, IndexError):
                pass

    if suspicious:
        for s in suspicious:
            add_log(f"⚠ {s}", "warn")
            add_threat("Suspicious.NetworkPort", s, "MEDIUM")
    else:
        add_log("Network connections: No suspicious ports detected.", "ok")

    lines = [l for l in out.splitlines() if "LISTENING" in l or "ESTABLISHED" in l]
    add_log(f"Active connections/listeners: {len(lines)} found.", "info")


def check_running_processes():
    """Check running processes for suspicious entries."""
    add_log("Scanning running processes...", "info")

    if platform.system() == "Windows":
        out = run_cmd("tasklist /v /fo CSV", "Running Processes")
        lines = out.splitlines()[1:]  # skip header
    else:
        out = run_cmd("ps aux", "Running Processes")
        lines = out.splitlines()[1:]

    suspicious_procs = []
    for line in lines:
        line_lower = line.lower()
        for pattern in SUSPICIOUS_PATTERNS:
            if re.search(pattern, line_lower):
                suspicious_procs.append(line.split(",")[0].strip('"'))
                break

    if suspicious_procs:
        for p in suspicious_procs[:10]:
            add_log(f"⚠ Suspicious process: {p}", "warn")
            add_threat("Suspicious.Process", f"Process: {p}", "MEDIUM")
    else:
        add_log(f"Processes scanned: {len(lines)}. No suspicious processes found.", "ok")


def check_startup_entries():
    """Check registry startup keys (Windows) or startup files."""
    add_log("Checking startup entries...", "info")

    if platform.system() == "Windows":
        keys = [
            r"HKLM\Software\Microsoft\Windows\CurrentVersion\Run",
            r"HKCU\Software\Microsoft\Windows\CurrentVersion\Run",
            r"HKLM\Software\Microsoft\Windows\CurrentVersion\RunOnce",
        ]
        for key in keys:
            out = run_cmd(f'reg query "{key}"', f"Registry: {key}")
            for line in out.splitlines():
                line_lower = line.lower()
                for pattern in SUSPICIOUS_PATTERNS:
                    if re.search(pattern, line_lower):
                        add_log(f"⚠ Suspicious startup entry: {line.strip()}", "warn")
                        add_threat("Suspicious.StartupEntry", line.strip(), "HIGH")
                        break
        add_log("Startup registry keys checked.", "ok")
    else:
        # Linux/Mac: check common startup locations
        startup_paths = [
            "/etc/init.d", "/etc/cron.d", os.path.expanduser("~/.bashrc"),
            os.path.expanduser("~/.profile"), "/etc/rc.local"
        ]
        for path in startup_paths:
            if os.path.exists(path):
                add_log(f"Startup path found: {path}", "info")
        add_log("Startup entries checked.", "ok")


def check_temp_directory():
    """Scan temp directory for executables and suspicious files."""
    add_log("Scanning TEMP directory...", "info")
    temp_dir = os.environ.get("TEMP", os.environ.get("TMPDIR", "/tmp"))

    if not os.path.exists(temp_dir):
        add_log("TEMP directory not found.", "warn")
        return

    count = 0
    for fname in os.listdir(temp_dir):
        fpath = os.path.join(temp_dir, fname)
        ext = os.path.splitext(fname)[1].lower()
        if ext in SUSPICIOUS_EXTENSIONS:
            count += 1
            add_log(f"⚠ Executable in TEMP: {fname}", "warn")
            add_threat("Suspicious.TempExecutable", fpath, "HIGH")
            scan_file_hash(fpath)

    if count == 0:
        add_log(f"TEMP directory clean. No executables found.", "ok")
    else:
        add_log(f"{count} suspicious file(s) found in TEMP.", "warn")


def check_dns_cache():
    """Check DNS cache for suspicious domains."""
    add_log("Checking DNS cache...", "info")
    suspicious_tlds = [".ru", ".cn", ".tk", ".xyz", ".top", ".pw"]

    if platform.system() == "Windows":
        out = run_cmd("ipconfig /displaydns", "DNS Cache")
    else:
        # Linux doesn't cache DNS by default; check /etc/hosts
        out = run_cmd("cat /etc/hosts", "Hosts File")

    flagged = []
    for line in out.splitlines():
        line_lower = line.lower()
        for tld in suspicious_tlds:
            if tld in line_lower and "windows" not in line_lower:
                flagged.append(line.strip())
                break

    if flagged:
        for f in flagged[:5]:
            add_log(f"⚠ Suspicious DNS entry: {f}", "warn")
        add_threat("Suspicious.DNSEntry", f"{len(flagged)} suspicious domain(s) in cache", "MEDIUM")
    else:
        add_log("DNS cache: No suspicious entries found.", "ok")


def check_services():
    """Check running services for suspicious entries."""
    add_log("Enumerating system services...", "info")

    if platform.system() == "Windows":
        out = run_cmd("sc query type= all state= all", "System Services")
    else:
        out = run_cmd("systemctl list-units --type=service --state=running", "System Services")

    service_count = out.count("SERVICE_NAME") if platform.system() == "Windows" else out.count(".service")
    add_log(f"Services found: {service_count}. Analyzing...", "info")

    for line in out.splitlines():
        line_lower = line.lower()
        for pattern in SUSPICIOUS_PATTERNS:
            if re.search(pattern, line_lower):
                add_log(f"⚠ Suspicious service: {line.strip()}", "warn")
                add_threat("Suspicious.Service", line.strip(), "HIGH")
                break

    add_log("Service analysis complete.", "ok")


def check_scheduled_tasks():
    """Check scheduled tasks for suspicious entries."""
    add_log("Checking scheduled tasks...", "info")

    if platform.system() == "Windows":
        out = run_cmd("schtasks /query /fo LIST /v", "Scheduled Tasks", timeout=20)
    else:
        out = run_cmd("crontab -l 2>/dev/null; ls /etc/cron* 2>/dev/null", "Cron Jobs")

    for line in out.splitlines():
        line_lower = line.lower()
        for pattern in SUSPICIOUS_PATTERNS:
            if re.search(pattern, line_lower):
                add_log(f"⚠ Suspicious scheduled task: {line.strip()}", "warn")
                add_threat("Suspicious.ScheduledTask", line.strip(), "HIGH")
                break

    add_log("Scheduled tasks checked.", "ok")


def check_security_logs():
    """Check Windows Security Event Log for failed logins, etc."""
    add_log("Reviewing security event logs...", "info")

    if platform.system() == "Windows":
        out = run_cmd(
            'wevtutil qe Security /c:100 /f:text /q:"*[System[(EventID=4625 or EventID=4648 or EventID=4672)]]"',
            "Security Event Log",
            timeout=20
        )
        failed_logins = out.count("Event ID: 4625") + out.lower().count("4625")
        if failed_logins > 5:
            add_log(f"⚠ {failed_logins} failed login attempts found!", "warn")
            add_threat("Suspicious.BruteForce", f"{failed_logins} failed login attempts in security log", "MEDIUM")
        else:
            add_log(f"Security log: {failed_logins} failed login(s). Within normal range.", "ok")
    else:
        out = run_cmd("last -n 20", "Login History")
        add_log("Login history reviewed.", "ok")


def run_sfc():
    """Run System File Checker (Windows only)."""
    if platform.system() != "Windows":
        add_log("SFC skipped (Windows only).", "info")
        return
    add_log("Running System File Checker (sfc /scannow)... This may take a moment.", "info")
    out = run_cmd("sfc /scannow", "System File Checker", timeout=120)
    if "no integrity violations" in out.lower():
        add_log("SFC: No integrity violations found.", "ok")
    elif "repaired" in out.lower() or "found" in out.lower():
        add_log("⚠ SFC: System files were repaired or corrupted files found.", "warn")
        add_threat("SystemFile.Corruption", "Windows System File Checker detected file corruption", "MEDIUM")
    else:
        add_log("SFC: Check complete.", "ok")


def scan_file_hash(filepath):
    """Compute SHA-256 hash of a file and check against known bad hashes."""
    try:
        sha256 = hashlib.sha256()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(65536), b""):
                sha256.update(chunk)
        file_hash = sha256.hexdigest()
        if file_hash in KNOWN_BAD_HASHES:
            add_log(f"🔴 HASH MATCH: {filepath} [{file_hash[:16]}...]", "threat")
            add_threat("Malware.KnownBadHash", filepath, "HIGH", file_hash)
            return True
        return False
    except (PermissionError, OSError):
        return False


def add_threat(name, path, severity, extra=""):
    threat = {
        "name": name,
        "path": path,
        "severity": severity,
        "extra": extra,
        "time": time.strftime("%H:%M:%S"),
    }
    scan_state["threats"].append(threat)
    add_log(f"THREAT DETECTED: {name} @ {path}", "threat")


# ─────────────────────────────────────────────
#  FILE SYSTEM SCANNER
# ─────────────────────────────────────────────

def scan_directory(root_path, scan_type="quick"):
    """Recursively scan a directory for threats."""
    add_log(f"Scanning directory: {root_path}", "info")
    skip_dirs = {
        "windows", "system volume information", "$recycle.bin",
        "program files (x86)", "perflogs", "recovery",
        ".git", "node_modules", "__pycache__", ".cache"
    }

    for dirpath, dirnames, filenames in os.walk(root_path, topdown=True):
        if not scan_state["running"]:
            break

        # Skip known safe system dirs in quick scan
        if scan_type == "quick":
            dirnames[:] = [
                d for d in dirnames
                if d.lower() not in skip_dirs
            ]

        for filename in filenames:
            if not scan_state["running"]:
                return

            filepath = os.path.join(dirpath, filename)
            ext = os.path.splitext(filename)[1].lower()
            scan_state["files_scanned"] += 1

            # Check suspicious filename patterns
            fname_lower = filename.lower()
            for pattern in SUSPICIOUS_PATTERNS:
                if re.search(pattern, fname_lower):
                    add_log(f"⚠ Suspicious filename: {filename}", "warn")
                    add_threat("Suspicious.Filename", filepath, "MEDIUM")
                    break

            # Check file hash if it's an executable
            if ext in SUSPICIOUS_EXTENSIONS:
                try:
                    size = os.path.getsize(filepath)
                    if 0 < size < 50 * 1024 * 1024:  # < 50MB
                        scan_file_hash(filepath)
                except OSError:
                    pass


# ─────────────────────────────────────────────
#  MAIN SCAN ORCHESTRATOR
# ─────────────────────────────────────────────

def run_scan(scan_type="quick", custom_path=None):
    """Run the full scan in a background thread."""
    scan_state["running"] = True
    scan_state["complete"] = False
    scan_state["progress"] = 0
    scan_state["files_scanned"] = 0
    scan_state["threats"] = []
    scan_state["logs"] = []
    scan_state["cmds_done"] = []
    scan_state["scan_type"] = scan_type

    add_log("=== OMNIANTI SCAN ENGINE STARTED ===", "info")
    add_log(f"Platform: {platform.system()} {platform.release()}", "info")
    add_log(f"Scan type: {scan_type.upper()}", "info")
    add_log(f"Python: {sys.version.split()[0]}", "info")

    steps = [
        (check_network_connections, 10),
        (check_running_processes, 20),
        (check_startup_entries, 35),
        (check_temp_directory, 50),
        (check_dns_cache, 60),
        (check_services, 70),
        (check_scheduled_tasks, 78),
        (check_security_logs, 85),
        (run_sfc, 90),
    ]

    for func, prog in steps:
        if not scan_state["running"]:
            break
        func()
        scan_state["progress"] = prog
        time.sleep(0.3)

    # File system scan
    if scan_state["running"]:
        add_log("Starting filesystem scan...", "info")
        scan_state["progress"] = 92

        if custom_path and os.path.exists(custom_path):
            scan_paths = [custom_path]
        elif scan_type == "quick":
            # Quick scan: only high-risk folders
            scan_paths = [
                os.environ.get("TEMP", "/tmp"),
                os.path.expanduser("~/Downloads"),
                os.path.expanduser("~/Desktop"),
            ]
            if platform.system() == "Windows":
                scan_paths += [
                    r"C:\Windows\Temp",
                    r"C:\Users\Public\Downloads",
                ]
        elif scan_type == "full":
            scan_paths = ["C:\\"] if platform.system() == "Windows" else ["/home", "/tmp", "/var"]
        else:
            scan_paths = [os.path.expanduser("~")]

        for path in scan_paths:
            if os.path.exists(path) and scan_state["running"]:
                scan_directory(path, scan_type)

    scan_state["progress"] = 100
    scan_state["running"] = False
    scan_state["complete"] = True

    threat_count = len(scan_state["threats"])
    add_log(f"=== SCAN COMPLETE: {scan_state['files_scanned']:,} files scanned ===", "info")
    if threat_count == 0:
        add_log("✅ No threats detected. System appears clean.", "ok")
    else:
        add_log(f"⚠ {threat_count} threat(s) found! Review and take action.", "threat")


# ─────────────────────────────────────────────
#  HTTP SERVER
# ─────────────────────────────────────────────

class OmniAntiHandler(BaseHTTPRequestHandler):

    def log_message(self, format, *args):
        pass  # Suppress default HTTP logs

    def send_json(self, data, status=200):
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Content-Length", len(body))
        self.end_headers()
        self.wfile.write(body)

    def serve_file(self, filepath, content_type):
        try:
            with open(filepath, "rb") as f:
                data = f.read()
            self.send_response(200)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", len(data))
            self.end_headers()
            self.wfile.write(data)
        except FileNotFoundError:
            self.send_response(404)
            self.end_headers()

    def do_OPTIONS(self):
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()

    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path
        params = parse_qs(parsed.query)

        if path == "/" or path == "/index.html":
            self.serve_file(os.path.join(os.path.dirname(__file__), "templates", "index.html"), "text/html")

        elif path == "/api/status":
            self.send_json({
                "running": scan_state["running"],
                "progress": scan_state["progress"],
                "files_scanned": scan_state["files_scanned"],
                "threat_count": len(scan_state["threats"]),
                "complete": scan_state["complete"],
                "scan_type": scan_state["scan_type"],
                "platform": platform.system(),
            })

        elif path == "/api/logs":
            since = int(params.get("since", ["0"])[0])
            self.send_json(scan_state["logs"][since:])

        elif path == "/api/threats":
            self.send_json(scan_state["threats"])

        elif path == "/api/cmds":
            self.send_json(scan_state["cmds_done"])

        elif path == "/api/start":
            if not scan_state["running"]:
                scan_type = params.get("type", ["quick"])[0]
                custom_path = params.get("path", [None])[0]
                t = threading.Thread(target=run_scan, args=(scan_type, custom_path), daemon=True)
                t.start()
                self.send_json({"status": "started", "scan_type": scan_type})
            else:
                self.send_json({"status": "already_running"})

        elif path == "/api/stop":
            scan_state["running"] = False
            self.send_json({"status": "stopped"})

        elif path == "/api/sysinfo":
            info = {
                "os": f"{platform.system()} {platform.release()}",
                "machine": platform.machine(),
                "processor": platform.processor(),
                "python": sys.version.split()[0],
                "hostname": platform.node(),
            }
            self.send_json(info)

        else:
            self.send_response(404)
            self.end_headers()


# ─────────────────────────────────────────────
#  ENTRY POINT
# ─────────────────────────────────────────────

if __name__ == "__main__":
    PORT = 5000
    print(f"""
╔══════════════════════════════════════════╗
║     OMNIANTI - Advanced Threat Protection    ║
║     Real System Scanner v1.0              ║
╠══════════════════════════════════════════╣
║  Server starting on http://localhost:{PORT}  ║
║  Press Ctrl+C to stop                    ║
╚══════════════════════════════════════════╝
    """)
    server = HTTPServer(("localhost", PORT), OmniAntiHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[OmniAnti] Server stopped.")
        server.server_close()
