#!/usr/bin/python3
"""
SetecMITM — generic IoT / cloud-device MITM framework (PyQt6 GUI).
Run: sudo /usr/bin/python3 gui.py
"""

import os
import sys
import json
import threading
import signal

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from PyQt6.QtCore import Qt, QTimer, pyqtSignal, QObject
from PyQt6.QtGui import QFont, QTextCursor, QColor
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QTabWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QPlainTextEdit, QLabel, QLineEdit, QTableWidget,
    QTableWidgetItem, QHeaderView, QFormLayout, QGroupBox,
    QStatusBar, QMessageBox, QTextEdit, QAbstractItemView, QCheckBox,
)

from config import Config, CONFIG_FILE
from utils.log import (
    log, log_lines, init_logfile, close_logfile, lock,
    C_NONE, C_ERROR, C_SUCCESS, C_INFO, C_TRAFFIC, C_IMPORTANT,
)
from utils import proto as proto_id
from services import intruder_watch
from inject import packet
from mitm import Controller


QT_COLORS = {
    C_NONE:      QColor("#cccccc"),
    C_ERROR:     QColor("#ff5555"),
    C_SUCCESS:   QColor("#50fa7b"),
    C_INFO:      QColor("#8be9fd"),
    C_TRAFFIC:   QColor("#f1fa8c"),
    C_IMPORTANT: QColor("#ff79c6"),
}


class LogBridge(QObject):
    new_lines = pyqtSignal(list)


class MainWindow(QMainWindow):
    def __init__(self, ctrl):
        super().__init__()
        self.ctrl = ctrl
        self.setWindowTitle("SetecMITM — Generic IoT MITM")
        self.resize(1400, 900)
        self._apply_dark_theme()

        self.bridge = LogBridge()
        self.bridge.new_lines.connect(self._append_log)
        self._last_log_idx = 0

        self.tabs = QTabWidget()
        self.tabs.addTab(self._build_dashboard(),  "Dashboard")
        self.tabs.addTab(self._build_log_tab(),    "Live Log")
        self.tabs.addTab(self._build_intruder_tab(), "Intruders")
        self.tabs.addTab(self._build_inject_tab(), "Inject")
        self.tabs.addTab(self._build_settings_tab(), "Settings")
        self.tabs.addTab(self._build_help_tab(),   "Help")
        self.setCentralWidget(self.tabs)

        self.status = QStatusBar()
        self.setStatusBar(self.status)
        self.refresh_timer = QTimer(self)
        self.refresh_timer.timeout.connect(self._tick)
        self.refresh_timer.start(300)

        log("SetecMITM GUI ready", C_SUCCESS)
        if not self.ctrl.cfg["target_ip"]:
            log("⚠ target_ip is not set — open the Settings tab first", C_ERROR)

    def _apply_dark_theme(self):
        self.setStyleSheet("""
            QWidget { background: #1e1f29; color: #f8f8f2; font-family: 'JetBrains Mono','Fira Code',monospace; font-size: 11pt; }
            QTabWidget::pane { border: 1px solid #44475a; }
            QTabBar::tab { background: #282a36; padding: 8px 16px; border: 1px solid #44475a; }
            QTabBar::tab:selected { background: #44475a; color: #50fa7b; }
            QPushButton { background: #44475a; border: 1px solid #6272a4; padding: 6px 14px; border-radius: 3px; }
            QPushButton:hover { background: #6272a4; }
            QPushButton:pressed { background: #50fa7b; color: #282a36; }
            QPlainTextEdit, QTextEdit, QLineEdit { background: #282a36; border: 1px solid #44475a; selection-background-color: #44475a; }
            QHeaderView::section { background: #44475a; color: #f8f8f2; padding: 4px; border: none; }
            QTableWidget { background: #282a36; gridline-color: #44475a; }
            QTableWidget::item:selected { background: #44475a; }
            QGroupBox { border: 1px solid #44475a; margin-top: 10px; padding-top: 10px; }
            QGroupBox::title { subcontrol-origin: margin; left: 10px; padding: 0 5px; color: #8be9fd; }
            QStatusBar { background: #282a36; color: #50fa7b; }
        """)

    # ── Dashboard ─────────────────────────────────────────
    def _build_dashboard(self):
        w = QWidget()
        layout = QVBoxLayout(w)

        ctrl_box = QGroupBox("MITM Control")
        cl = QHBoxLayout(ctrl_box)
        for label, fn in [
            ("▶ START ALL", lambda: threading.Thread(target=self.ctrl.start_services, daemon=True).start()),
            ("⏹ STOP ALL",  lambda: threading.Thread(target=self.ctrl.stop_services,  daemon=True).start()),
            ("Clear Log",   self._clear_log),
        ]:
            b = QPushButton(label); b.clicked.connect(fn); cl.addWidget(b)
        cl.addStretch()
        layout.addWidget(ctrl_box)

        self.state_label = QLabel("MITM: STOPPED")
        self.state_label.setStyleSheet("color:#ff5555; font-size:18pt; font-weight:bold; padding:10px;")
        layout.addWidget(self.state_label)

        flags_box = QGroupBox("Services (click to toggle)")
        fl = QVBoxLayout(flags_box)
        self.svc_buttons = {}
        for name in ("arp", "dns", "http", "https", "sniffer", "intruder"):
            btn = QPushButton(f"● {name}: off")
            btn.setStyleSheet("text-align:left; color:#ff5555; padding:6px; background:#282a36;")
            btn.clicked.connect(lambda _, n=name: threading.Thread(
                target=self.ctrl.toggle_service, args=(n,), daemon=True).start())
            self.svc_buttons[name] = btn
            fl.addWidget(btn)
        layout.addWidget(flags_box)

        proto_box = QGroupBox("Protocols Seen")
        pl = QVBoxLayout(proto_box)
        self.proto_label = QLabel("(none yet)")
        self.proto_label.setStyleSheet("color:#f1fa8c; font-family:monospace;")
        pl.addWidget(self.proto_label)
        layout.addWidget(proto_box)

        info_box = QGroupBox("Target")
        il = QFormLayout(info_box)
        self.lbl_tgt = QLabel(self.ctrl.cfg["target_ip"] or "(unset)")
        self.lbl_us  = QLabel(self.ctrl.cfg["our_ip"]    or "(unset)")
        self.lbl_rtr = QLabel(self.ctrl.cfg["router_ip"] or "(unset)")
        self.lbl_mac = QLabel(self.ctrl.cfg["target_mac"] or "(unset)")
        for lbl in (self.lbl_tgt, self.lbl_us, self.lbl_rtr, self.lbl_mac):
            lbl.setStyleSheet("color:#f1fa8c; font-weight:bold;")
        il.addRow("Target IP:", self.lbl_tgt)
        il.addRow("Our IP:",    self.lbl_us)
        il.addRow("Router IP:", self.lbl_rtr)
        il.addRow("Target MAC:", self.lbl_mac)
        layout.addWidget(info_box)

        layout.addStretch()
        return w

    # ── Live Log ──────────────────────────────────────────
    def _build_log_tab(self):
        w = QWidget()
        layout = QVBoxLayout(w)
        bar = QHBoxLayout()
        bar.addWidget(QLabel("Filter:"))
        self.log_filter = QLineEdit()
        self.log_filter.setPlaceholderText("substring filter (live)…")
        bar.addWidget(self.log_filter)
        self.autoscroll_cb = QCheckBox("Autoscroll")
        self.autoscroll_cb.setChecked(True)
        bar.addWidget(self.autoscroll_cb)
        b = QPushButton("Clear")
        b.clicked.connect(self._clear_log)
        bar.addWidget(b)
        layout.addLayout(bar)
        self.log_view = QTextEdit()
        self.log_view.setReadOnly(True)
        self.log_view.setLineWrapMode(QTextEdit.LineWrapMode.NoWrap)
        f = QFont("JetBrains Mono", 10)
        f.setStyleHint(QFont.StyleHint.Monospace)
        self.log_view.setFont(f)
        layout.addWidget(self.log_view)
        return w

    # ── Intruders ─────────────────────────────────────────
    def _build_intruder_tab(self):
        w = QWidget()
        layout = QVBoxLayout(w)
        head = QHBoxLayout()
        self.intruder_count = QLabel("0 events")
        self.intruder_count.setStyleSheet("color:#ff79c6; font-size:14pt; font-weight:bold;")
        head.addWidget(self.intruder_count)
        head.addStretch()
        b = QPushButton("Clear")
        b.clicked.connect(lambda: (intruder_watch.clear_intruders(), self._refresh_intruders()))
        head.addWidget(b)
        layout.addLayout(head)
        self.intruder_table = QTableWidget(0, 5)
        self.intruder_table.setHorizontalHeaderLabels(["Time", "Kind", "Source", "Destination", "Detail"])
        self.intruder_table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Stretch)
        self.intruder_table.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        self.intruder_table.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)
        layout.addWidget(self.intruder_table)
        return w

    # ── Inject ────────────────────────────────────────────
    def _build_inject_tab(self):
        w = QWidget()
        layout = QVBoxLayout(w)

        udp_box = QGroupBox("UDP Inject")
        ul = QFormLayout(udp_box)
        self.udp_ip = QLineEdit(self.ctrl.cfg["target_ip"])
        self.udp_port = QLineEdit("0")
        self.udp_payload = QLineEdit()
        self.udp_payload.setPlaceholderText("hex payload, e.g. deadbeef")
        ul.addRow("Dst IP:", self.udp_ip)
        ul.addRow("Dst port:", self.udp_port)
        ul.addRow("Payload:", self.udp_payload)
        b = QPushButton("Send UDP")
        b.clicked.connect(self._send_udp)
        ul.addRow("", b)
        layout.addWidget(udp_box)

        arp_box = QGroupBox("ARP Reply")
        al = QFormLayout(arp_box)
        self.arp_src = QLineEdit(self.ctrl.cfg["router_ip"])
        self.arp_dst = QLineEdit(self.ctrl.cfg["target_ip"])
        al.addRow("Src IP (spoof):", self.arp_src)
        al.addRow("Dst IP:", self.arp_dst)
        b2 = QPushButton("Send ARP")
        b2.clicked.connect(lambda: packet.inject(self.ctrl.cfg, {
            "type": "arp_reply", "src_ip": self.arp_src.text(), "dst_ip": self.arp_dst.text()
        }))
        al.addRow("", b2)
        layout.addWidget(arp_box)

        dns_box = QGroupBox("DNS Query")
        dl = QFormLayout(dns_box)
        self.dns_dom = QLineEdit()
        self.dns_dom.setPlaceholderText("example.com")
        dl.addRow("Domain:", self.dns_dom)
        b3 = QPushButton("Send DNS")
        b3.clicked.connect(lambda: packet.inject(self.ctrl.cfg, {
            "type": "dns_query", "domain": self.dns_dom.text()
        }))
        dl.addRow("", b3)
        layout.addWidget(dns_box)

        layout.addStretch()
        return w

    def _send_udp(self):
        try:
            packet.inject(self.ctrl.cfg, {
                "type": "udp",
                "dst_ip": self.udp_ip.text(),
                "dst_port": int(self.udp_port.text()),
                "payload": self.udp_payload.text(),
                "payload_hex": True,
            })
        except Exception as e:
            QMessageBox.warning(self, "Inject error", str(e))

    # ── Settings ──────────────────────────────────────────
    def _build_settings_tab(self):
        w = QWidget()
        outer = QVBoxLayout(w)

        info = QLabel(
            "Set the target's IP, the target's MAC, the IP of THIS box, the gateway, and the network interface.\n"
            "Save Config to persist. Restart MITM after changing target_ip."
        )
        info.setStyleSheet("color:#8be9fd; padding:6px;")
        outer.addWidget(info)

        form = QFormLayout()
        self.cfg_inputs = {}
        for k, v in self.ctrl.cfg.items():
            if k.startswith("_"):
                continue
            le = QLineEdit(json.dumps(v) if not isinstance(v, str) else v)
            self.cfg_inputs[k] = le
            form.addRow(k, le)
        outer.addLayout(form)

        bb = QHBoxLayout()
        b1 = QPushButton("Save Config")
        b1.clicked.connect(self._save_config)
        bb.addWidget(b1)
        b2 = QPushButton("Reload From Disk")
        b2.clicked.connect(self._reload_config)
        bb.addWidget(b2)
        bb.addStretch()
        path_lbl = QLabel(f"file: {CONFIG_FILE}")
        path_lbl.setStyleSheet("color:#888; font-size:10pt;")
        bb.addWidget(path_lbl)
        outer.addLayout(bb)
        return w

    def _save_config(self):
        for k, le in self.cfg_inputs.items():
            old = self.ctrl.cfg[k]
            v = le.text()
            try:
                if isinstance(old, bool):
                    v = v.lower() in ("true", "1", "yes")
                elif isinstance(old, int):
                    v = int(v)
                elif isinstance(old, float):
                    v = float(v)
                elif isinstance(old, list):
                    v = json.loads(v)
            except (ValueError, json.JSONDecodeError):
                pass
            self.ctrl.cfg[k] = v
        self.ctrl.cfg.save()
        log("config saved from GUI", C_SUCCESS)

    def _reload_config(self):
        self.ctrl.cfg.load()
        for k, le in self.cfg_inputs.items():
            v = self.ctrl.cfg[k]
            le.setText(json.dumps(v) if not isinstance(v, str) else v)
        log("config reloaded from disk", C_SUCCESS)

    # ── Help ──────────────────────────────────────────────
    def _build_help_tab(self):
        w = QWidget()
        layout = QVBoxLayout(w)
        view = QTextEdit()
        view.setReadOnly(True)
        view.setHtml("""
<h2 style="color:#50fa7b">SetecMITM</h2>
<p>Generic LAN-side MITM framework for any IoT or cloud-connected device.
Built for authorized security research on hardware you own.</p>

<h3 style="color:#8be9fd">Quick start</h3>
<ol>
<li>Open the <b>Settings</b> tab. Fill in <code>target_ip</code>, <code>target_mac</code>, <code>our_ip</code>, <code>router_ip</code>, and <code>iface</code>. Save.</li>
<li>Open the <b>Dashboard</b> tab.</li>
<li>Click <b>▶ START ALL</b> — or click each service row individually.</li>
<li>Switch to <b>Live Log</b> to watch traffic in real time.</li>
<li>Switch to <b>Intruders</b> to see detected suspicious activity.</li>
</ol>

<h3 style="color:#8be9fd">Tabs</h3>
<ul>
<li><b>Dashboard</b> — START/STOP, click any service to toggle, watch protocol counts and target info.</li>
<li><b>Live Log</b> — every log line, color-coded. Filter by substring. Toggle Autoscroll.</li>
<li><b>Intruders</b> — table of ARP-spoof attempts, unknown LAN peers contacting the target, and outbound destinations not on your whitelist.</li>
<li><b>Inject</b> — craft and send raw UDP, ARP, or DNS packets.</li>
<li><b>Settings</b> — every config key, editable, persisted to <code>~/.config/setec-mitm/config.json</code>.</li>
</ul>

<h3 style="color:#8be9fd">Services</h3>
<ul>
<li><b>arp</b> — ARP cache poisoning so the target thinks we are the gateway.</li>
<li><b>dns</b> — DNS spoof: redirect cloud lookups to our box.</li>
<li><b>http / https</b> — Intercept ports 80/443. HTTPS uses an auto-generated cert with full SAN list (regen with <code>regen_cert.sh</code>).</li>
<li><b>sniffer</b> — Raw packet sniffer with conntrack original-destination lookup and protocol fingerprinting.</li>
<li><b>intruder</b> — Detects ARP spoofs against the target, unknown LAN peers contacting it, and outbound destinations not in <code>intruder_known_nets</code>.</li>
</ul>

<h3 style="color:#8be9fd">Plugins (target-specific code)</h3>
<p>Vendor-specific clients (cloud API, fuzz wordlists, CVE checks) live under
<code>targets/&lt;name&gt;/plugin.py</code>. Set <code>target_plugin</code> in
the Settings tab to load one. See <code>targets/example/</code> for the layout.</p>

<h3 style="color:#8be9fd">Run</h3>
<p><code>sudo /usr/bin/python3 gui.py</code></p>
        """)
        layout.addWidget(view)
        return w

    # ── Periodic refresh ──────────────────────────────────
    def _tick(self):
        with lock:
            total = len(log_lines)
            if total < self._last_log_idx:
                self._last_log_idx = 0
            new = list(log_lines)[self._last_log_idx:]
            self._last_log_idx = total
        if new:
            self.bridge.new_lines.emit(new)

        if self.ctrl.services_running:
            self.state_label.setText("MITM: RUNNING")
            self.state_label.setStyleSheet("color:#50fa7b; font-size:18pt; font-weight:bold; padding:10px;")
        else:
            self.state_label.setText("MITM: STOPPED")
            self.state_label.setStyleSheet("color:#ff5555; font-size:18pt; font-weight:bold; padding:10px;")

        for name, btn in self.svc_buttons.items():
            on = self.ctrl.flags.get(name, False)
            btn.setText(f"● {name}: {'ON' if on else 'off'}")
            color = "#50fa7b" if on else "#ff5555"
            btn.setStyleSheet(f"text-align:left; color:{color}; padding:6px; background:#282a36;")

        counts = proto_id.seen_counts()
        if counts:
            txt = "  ".join(f"{k}={v}" for k, v in sorted(counts.items(), key=lambda x: -x[1]))
            self.proto_label.setText(txt)

        self._refresh_intruders()

        self.lbl_tgt.setText(self.ctrl.cfg["target_ip"] or "(unset)")
        self.lbl_us.setText(self.ctrl.cfg["our_ip"] or "(unset)")
        self.lbl_rtr.setText(self.ctrl.cfg["router_ip"] or "(unset)")
        self.lbl_mac.setText(self.ctrl.cfg["target_mac"] or "(unset)")

        self.status.showMessage(
            f"target={self.ctrl.cfg['target_ip'] or '?'}  iface={self.ctrl.cfg['iface'] or '?'}  "
            f"intruders={len(intruder_watch.get_intruders())}"
        )

    def _refresh_intruders(self):
        items = intruder_watch.get_intruders()
        self.intruder_count.setText(f"{len(items)} events")
        if self.intruder_table.rowCount() != len(items):
            self.intruder_table.setRowCount(len(items))
            for i, e in enumerate(items):
                self.intruder_table.setItem(i, 0, QTableWidgetItem(e["ts"]))
                kind_item = QTableWidgetItem(e["kind"])
                kind_item.setForeground(QColor("#ff79c6"))
                self.intruder_table.setItem(i, 1, kind_item)
                self.intruder_table.setItem(i, 2, QTableWidgetItem(e["src"]))
                self.intruder_table.setItem(i, 3, QTableWidgetItem(e["dst"]))
                self.intruder_table.setItem(i, 4, QTableWidgetItem(e["detail"]))

    def _append_log(self, lines):
        flt = self.log_filter.text().lower()
        autoscroll = self.autoscroll_cb.isChecked()
        sb = self.log_view.verticalScrollBar()
        old_pos = sb.value()
        old_user_cursor = self.log_view.textCursor()
        write_cursor = QTextCursor(self.log_view.document())
        write_cursor.movePosition(QTextCursor.MoveOperation.End)
        for line, color in lines:
            if flt and flt not in line.lower():
                continue
            fmt = write_cursor.charFormat()
            fmt.setForeground(QT_COLORS.get(color, QT_COLORS[C_NONE]))
            write_cursor.setCharFormat(fmt)
            write_cursor.insertText(line + "\n")
        if autoscroll:
            self.log_view.moveCursor(QTextCursor.MoveOperation.End)
            self.log_view.ensureCursorVisible()
        else:
            self.log_view.setTextCursor(old_user_cursor)
            sb.setValue(old_pos)

    def _clear_log(self):
        self.log_view.clear()
        with lock:
            log_lines.clear()
        self._last_log_idx = 0

    def closeEvent(self, ev):
        if self.ctrl.services_running:
            self.ctrl.stop_services()
        self.ctrl.running = False
        close_logfile()
        ev.accept()


def main():
    if os.geteuid() != 0:
        print("Run with: sudo /usr/bin/python3 gui.py")
        sys.exit(1)
    ctrl = Controller()
    os.makedirs(ctrl.cfg["log_dir"], exist_ok=True)
    init_logfile(f"{ctrl.cfg['log_dir']}/setec_mitm.log")
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    app = QApplication(sys.argv)
    win = MainWindow(ctrl)
    win.show()
    rc = app.exec()
    if ctrl.services_running:
        ctrl.stop_services()
    close_logfile()
    sys.exit(rc)


if __name__ == "__main__":
    main()