#!/usr/bin/env python3
"""termmap - an ncurses front-end for nmap.

Runs nmap scans and streams the output, with the common scan types on a
toolbar: quick, deep, ping, ports, service/version, OS, UDP and scripts.
"""

import curses
import os
import queue
import shlex
import subprocess
import sys
import textwrap
import threading
import time

VERSION = "0.1"
NMAP = "/usr/bin/nmap"


class Scan:
    def __init__(self):
        self.proc = None
        self.lines = []
        self.q = queue.Queue()
        self.running = False
        self.command = ""

    def start(self, args, target):
        self.stop()
        self.lines = []
        self.command = " ".join(["stdbuf", "-oL", "-eL", NMAP] + args + [target])
        self.lines.append("$ " + self.command)
        self.lines.append("")
        try:
            self.proc = subprocess.Popen(
                ["stdbuf", "-oL", "-eL", NMAP] + args + [target],
                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                text=True, errors="replace", bufsize=1, start_new_session=True,
            )
        except Exception as exc:  # noqa: BLE001
            self.lines.append("Could not start nmap: %s" % exc)
            return
        self.running = True
        threading.Thread(target=self._reader, args=(self.proc,), daemon=True).start()

    def _reader(self, proc):
        try:
            for line in proc.stdout:
                clean = "".join(ch if (ch.isprintable() or ch == "\t") else " " for ch in line.rstrip("\n"))
                self.q.put(clean)
        finally:
            rc = proc.wait()
            self.q.put(("__done__", rc))

    def pump(self):
        finished = None
        while True:
            try:
                item = self.q.get_nowait()
            except queue.Empty:
                break
            if isinstance(item, tuple):
                finished = item[1]
            else:
                self.lines.append(item)
        if finished is not None:
            self.running = False
            self.lines.append("")
            self.lines.append("=== nmap finished (exit %s) ===" % finished)
        return finished

    def stop(self):
        if self.proc and self.proc.poll() is None:
            try:
                self.proc.terminate()
                self.proc.wait(timeout=2)
            except Exception:
                try:
                    self.proc.kill()
                except Exception:
                    pass
        self.running = False


class App:
    def __init__(self, stdscr):
        self.scr = stdscr
        self.color = curses.has_colors()
        self.scan = Scan()
        self.target = "127.0.0.1"
        self.ports = ""
        self.options = ""
        self.scroll = 0
        self.auto_bottom = True
        self.focus = "output"
        self.button = 0
        self.message = ""
        self.version = "?"
        self.buttons = [
            ("Quick", "quick"),
            ("Deep", "deep"),
            ("Ping", "ping"),
            ("Ports", "ports"),
            ("Svc", "svc"),
            ("OS", "os"),
            ("UDP", "udp"),
            ("Scripts", "scripts"),
            ("Target", "target"),
            ("Opts", "opts"),
            ("Save", "save"),
            ("Quit", "quit"),
        ]
        self.button_xs = []
        self.width = 80

    # ---------- drawing ----------
    def size(self):
        return self.scr.getmaxyx()

    def draw_bar(self, y, text, focus=False, right=""):
        _, w = self.size()
        attr = (curses.color_pair(1) | (curses.A_BOLD if focus else 0)) if self.color else curses.A_REVERSE
        avail = max(0, w - 1 - (len(right) + 1 if right else 0))
        self.scr.addnstr(y, 0, text[:avail].ljust(avail), avail, attr)
        if right:
            self.scr.addnstr(y, w - 1 - len(right), right, len(right), attr | curses.A_BOLD)

    def draw_buttons(self):
        _, w = self.size()
        x = 0
        self.button_xs = []
        base = curses.color_pair(2) if self.color else curses.A_DIM
        self.scr.addnstr(1, 0, " " * (w - 1), w - 1, base)
        for i, (label, _) in enumerate(self.buttons):
            if x >= w - 2:
                break
            self.button_xs.append(x)
            chunk = label + " "
            if i == self.button and self.focus == "buttons":
                attr = curses.color_pair(3) if self.color else curses.A_REVERSE
            else:
                attr = base
            try:
                self.scr.addnstr(1, x, chunk, max(0, w - 1 - x), attr)
            except curses.error:
                pass
            x += len(chunk)

    def draw_info(self):
        h, w = self.size()
        info = " Target: %s   Ports: %s   Options: %s" % (
            self.target, self.ports or "(default)", self.options or "-")
        self.scr.addnstr(2, 0, info[: w - 1], w - 1, curses.A_DIM)

    def wrapped(self):
        width = max(20, self.size()[1] - 2)
        out = []
        for line in self.scan.lines:
            if not line:
                out.append("")
                continue
            parts = textwrap.wrap(line, width=width, replace_whitespace=False, drop_whitespace=False)
            out.extend(parts or [""])
        return out

    def draw_output(self):
        h, w = self.size()
        self.scr.move(3, 0)
        self.scr.clrtobot()
        self._wrapped = self.wrapped()
        visible = max(1, h - 5)
        if self.auto_bottom:
            self.scroll = max(0, len(self._wrapped) - visible)
        for i in range(visible):
            idx = self.scroll + i
            if idx >= len(self._wrapped):
                break
            self.scr.addnstr(3 + i, 0, self._wrapped[idx][: w - 1], w - 1)
        self._visible = visible

    def draw_status(self):
        h, w = self.size()
        state = "scanning" if self.scan.running else "idle"
        focus = {"output": "output", "buttons": "buttons"}[self.focus]
        msg = self.message or "Tab: %s | Enter: scan | t: target | p: ports | x: options | Esc: stop | q: quit" % focus
        status = " %s  [%s]  %d lines" % (msg, state, len(self.scan.lines))
        self.scr.move(h - 1, 0)
        self.scr.clrtoeol()
        attr = curses.color_pair(1) if self.color else curses.A_REVERSE
        self.scr.addnstr(h - 1, 0, status.ljust(w - 1), w - 1, attr)
        self.message = ""

    def redraw(self):
        self.scr.erase()
        right = " running " if self.scan.running else " idle "
        self.draw_bar(0, " nmap %s  -  termmap" % self.version, True, right)
        self.draw_buttons()
        self.draw_info()
        self.draw_output()
        self.draw_status()
        self.scr.refresh()

    # ---------- dialogs ----------
    def input_box(self, title, initial=""):
        h, w = self.size()
        win_w = min(w - 4, 70)
        win = curses.newwin(5, win_w, max(1, h // 2 - 2), max(1, (w - win_w) // 2))
        win.keypad(True)
        text = list(initial)
        pos = len(text)
        first = True
        while True:
            win.erase()
            win.box()
            win.addnstr(0, 2, " %s " % title, win_w - 4, curses.A_BOLD)
            win.addnstr(2, 2, "".join(text)[-(win_w - 4):], win_w - 4)
            win.addnstr(3, 2, "Enter: ok   Esc: cancel", win_w - 4, curses.A_DIM)
            win.move(2, 2 + min(pos, win_w - 5))
            win.refresh()
            key = win.getch()
            if key in (10, 13, curses.KEY_ENTER):
                return "".join(text)
            if key == 27:
                return None
            if key in (curses.KEY_BACKSPACE, 127, 8):
                if first:
                    text = []
                    first = False
                elif text:
                    text.pop()
                pos = len(text)
            elif key == 21:
                text = []
                first = False
                pos = 0
            elif 32 <= key < 127:
                if first:
                    text = []
                    first = False
                text.append(chr(key))
                pos = len(text)

    # ---------- scans ----------
    def base_args(self, extra):
        args = shlex.split(self.options) if self.options.strip() else []
        if self.ports.strip() and "-p" not in args:
            args += ["-p", self.ports.strip()]
        args += extra
        return args

    def run_scan(self, kind):
        if self.scan.running:
            self.message = "Already running; press Esc to stop"
            return
        tgt = self.target.strip()
        if not tgt:
            self.message = "Set a target first"
            return
        extra = {
            "quick": ["-T4", "-F"],
            "deep": ["-T4", "-A", "-v", "--version-light", "--host-timeout", "4m"],
            "ping": ["-sn"],
            "ports": ["-T4"],
            "svc": ["-sV", "--version-light", "--host-timeout", "3m"],
            "os": ["-O"],
            "udp": ["-sU", "--top-ports", "50"],
            "scripts": ["-sC"],
        }[kind]
        if kind in ("os", "udp") and os.geteuid() != 0:
            self.message = "OS/UDP scans need root: run 'sudo termmap'"
            return
        self.scan.start(self.base_args(extra), tgt)
        self.auto_bottom = True
        self.message = "Started %s scan" % kind

    def save(self):
        stamp = time.strftime("%Y%m%d-%H%M%S")
        path = os.path.expanduser("~/nmap-%s.txt" % stamp)
        try:
            with open(path, "w", encoding="utf-8") as fh:
                fh.write("\n".join(self.scan.lines) + "\n")
            self.message = "Saved to %s" % path
        except Exception as exc:  # noqa: BLE001
            self.message = "Save failed: %s" % exc

    def action(self, name):
        if name == "quit":
            self.scan.stop()
            raise SystemExit
        if name in ("quick", "deep", "ping", "ports", "svc", "os", "udp", "scripts"):
            self.run_scan(name)
        elif name == "target":
            t = self.input_box("Target (IP, hostname or CIDR)", self.target)
            if t is not None:
                self.target = t
        elif name == "ports":
            p = self.input_box("Ports (for example 22,80,443 or 1-1000)", self.ports)
            if p is not None:
                self.ports = p
        elif name == "opts":
            o = self.input_box("Extra nmap options", self.options)
            if o is not None:
                self.options = o
        elif name == "save":
            self.save()

    # ---------- input ----------
    def handle_key(self, key):
        if key == 9:
            self.focus = "buttons" if self.focus == "output" else "output"
            return
        if key == 27:
            if self.scan.running:
                self.scan.stop()
                self.message = "Scan stopped"
            return
        if key == ord("q"):
            self.scan.stop()
            raise SystemExit
        if key == ord("?"):
            self.show_help()
            return
        if key == ord("s"):
            self.save()
            return
        if key == ord("t"):
            self.action("target")
            return
        if key == ord("p"):
            self.action("ports")
            return
        if key == ord("x"):
            self.action("opts")
            return
        if self.focus == "buttons":
            if key in (curses.KEY_LEFT, ord("h")):
                self.button = (self.button - 1) % len(self.buttons)
            elif key in (curses.KEY_RIGHT, ord("l")):
                self.button = (self.button + 1) % len(self.buttons)
            elif key in (10, 13, curses.KEY_ENTER):
                self.action(self.buttons[self.button][1])
            return
        visible = getattr(self, "_visible", 10)
        if key in (curses.KEY_UP, ord("k")):
            self.auto_bottom = False
            self.scroll = max(0, self.scroll - 1)
        elif key in (curses.KEY_DOWN, ord("j")):
            self.scroll += 1
            if self.scroll >= max(0, len(getattr(self, "_wrapped", [])) - visible):
                self.auto_bottom = True
        elif key == curses.KEY_PPAGE:
            self.auto_bottom = False
            self.scroll = max(0, self.scroll - visible)
        elif key == curses.KEY_NPAGE:
            self.scroll += visible
        elif key == ord("g"):
            self.auto_bottom = False
            self.scroll = 0
        elif key == ord("G"):
            self.auto_bottom = True
        elif key == ord("t"):
            self.action("target")
        elif key == ord("p"):
            self.action("ports")
        elif key == ord("x"):
            self.action("opts")
        elif key == ord("s"):
            self.save()
        elif key == ord("?"):
            self.show_help()
        elif key == ord("q"):
            self.scan.stop()
            raise SystemExit

    def show_help(self):
        h, w = self.size()
        entries = [
            "Tab            switch between output and buttons",
            "Enter (button) run the selected scan, or open its dialog",
            "Up/Down        scroll the output",
            "g / G          top / bottom",
            "t / p / x      edit target, ports, extra options",
            "s              save the output to ~/nmap-<time>.txt",
            "Esc            stop a running scan",
            "q              quit",
        ]
        win_h = min(len(entries) + 3, h - 5)
        win_w = min(w - 4, 74)
        win = curses.newwin(win_h, win_w, max(2, (h - win_h) // 2), max(1, (w - win_w) // 2))
        win.keypad(True)
        win.erase()
        win.box()
        win.addnstr(0, 2, " Help ", win_w - 4, curses.A_BOLD)
        for i, line in enumerate(entries[: win_h - 2]):
            win.addnstr(1 + i, 2, line[: win_w - 4], win_w - 4)
        win.refresh()
        win.getch()

    def run(self):
        try:
            out = subprocess.run([NMAP, "--version"], capture_output=True, text=True, timeout=5)
            first = (out.stdout or "").splitlines()
            for line in first:
                if line.startswith("Nmap version"):
                    self.version = line.split()[2]
                    break
        except Exception:
            self.version = "?"
        self.scr.timeout(200)
        while True:
            self.scan.pump()
            self.redraw()
            key = self.scr.getch()
            if key == -1:
                continue
            if key == curses.KEY_RESIZE:
                continue
            if key == curses.KEY_MOUSE:
                try:
                    _, mx, my, _, bstate = curses.getmouse()
                except curses.error:
                    continue
                if bstate & curses.BUTTON1_CLICKED:
                    if my == 1:
                        idx = 0
                        for i, x in enumerate(self.button_xs):
                            if x <= mx:
                                idx = i
                        self.button = idx
                        self.focus = "buttons"
                        self.action(self.buttons[idx][1])
                    else:
                        self.focus = "output"
                    continue
                if bstate & curses.BUTTON4_PRESSED:
                    self.auto_bottom = False
                    self.scroll = max(0, self.scroll - 3)
                elif bstate & curses.BUTTON5_PRESSED:
                    self.scroll += 3
                continue
            self.handle_key(key)


def main(stdscr):
    curses.curs_set(0)
    stdscr.keypad(True)
    try:
        curses.mousemask(curses.ALL_MOUSE_EVENTS)
    except curses.error:
        pass
    try:
        curses.start_color()
        curses.use_default_colors()
        curses.init_pair(1, curses.COLOR_WHITE, curses.COLOR_BLUE)
        curses.init_pair(2, curses.COLOR_BLACK, curses.COLOR_CYAN)
        curses.init_pair(3, curses.COLOR_BLACK, curses.COLOR_YELLOW)
    except curses.error:
        pass
    App(stdscr).run()


if __name__ == "__main__":
    if "--version" in sys.argv:
        print("termmap %s" % VERSION)
        raise SystemExit(0)
    if "-h" in sys.argv or "--help" in sys.argv:
        print("usage: termmap")
        print("An ncurses front-end for nmap with a toolbar of common scans.")
        raise SystemExit(0)
    if os.environ.get("TERM", "") in ("", "dumb", "unknown"):
        os.environ["TERM"] = "xterm-256color"
    try:
        curses.wrapper(main)
    except Exception:
        import traceback
        with open("/tmp/termmap.err", "w", encoding="utf-8") as fh:
            traceback.print_exc(file=fh)
        raise
