main.py

import asyncio
import re
import shlex
import subprocess

from textual import on
from textual.app import App, ComposeResult
from textual.containers import Vertical
from textual.widgets import Input, RichLog, Static


class COMP7904(App):
    BINDINGS = [("q", "quit_app", "Quit"), ("c", "clear_output", "Clear Output")]

    TOOLS = {
        "1": {
            "name": "nmap",
            "placeholder": "Enter nmap command (example: -sV scanme.nmap.org)",
        },
        "2": {
            "name": "tcpdump",
            "placeholder": "Enter tcpdump command (example: -i any -c 10)",
        },
    }

    CSS = """
    Screen {
        align: center middle;
    }

    #panel {
        width: 90%;
        max-width: 120;
        height: 90%;
        border: round #666666;
        padding: 1 2;
    }

    #menu {
        margin-bottom: 1;
    }

    #nmap_input {
        margin-bottom: 1;
    }

    #output {
        height: 1fr;
        border: round #444444;
    }
    """

    def compose(self) -> ComposeResult:
        menu = (
            "COMP7904 CLI\n"
            "Press a number key to choose a tool:\n"
            "  1: nmap\n"
            "  2: tcpdump\n"
            "  3..9: reserved\n"
            "Press c to clear output.\n"
            "Press q to quit."
        )
        with Vertical(id="panel"):
            yield Static(menu, id="menu")
            yield Input(
                placeholder="Enter nmap command (example: -sV scanme.nmap.org)",
                id="nmap_input",
            )
            yield RichLog(id="output", wrap=True, highlight=True)

    def on_mount(self) -> None:
        self.current_tool: str | None = None
        nmap_input = self.query_one("#nmap_input", Input)
        nmap_input.styles.display = "none"
        nmap_input.disabled = True
        self.set_focus(None)
        self._write_output("Ready. Press 1 for nmap, 2 for tcpdump, q to quit.")

    def action_quit_app(self) -> None:
        self.exit()

    def action_clear_output(self) -> None:
        output = self.query_one("#output", RichLog)
        output.clear()
        self._write_output(
            "Output cleared. Press 1 for nmap, 2 for tcpdump, q to quit."
        )

    async def on_key(self, event) -> None:
        if event.key in self.TOOLS:
            self._show_tool_prompt(event.key)
            event.stop()
            return

        if event.key.isdigit():
            self._write_output(f"Tool {event.key} is not wired yet.")
            event.stop()

    def _show_tool_prompt(self, tool_key: str) -> None:
        self.current_tool = self.TOOLS[tool_key]["name"]
        nmap_input = self.query_one("#nmap_input", Input)
        nmap_input.value = ""
        nmap_input.placeholder = self.TOOLS[tool_key]["placeholder"]
        nmap_input.disabled = False
        nmap_input.styles.display = "block"
        nmap_input.focus()
        self._write_output(
            f"Enter your {self.current_tool} command. You can type full command or only args."
        )

    @on(Input.Submitted, "#nmap_input")
    async def handle_nmap_submit(self, event: Input.Submitted) -> None:
        nmap_input = self.query_one("#nmap_input", Input)
        nmap_input.styles.display = "none"
        nmap_input.disabled = True
        self.set_focus(None)

        tool = self.current_tool
        if not tool:
            self._write_output("No tool selected. Press 1 for nmap or 2 for tcpdump.")
            return

        command_text = event.value.strip()
        if not command_text:
            self._write_output("No command entered. Press 1 or 2 to try again.")
            return

        try:
            command = self._normalize_tool_command(tool, command_text)
        except ValueError as exc:
            self._write_output(f"Invalid command: {exc}")
            return

        self._write_output("")
        self._write_output(f"Running: {' '.join(command)}")
        result = await self._run_command(tool, command)

        if result is None:
            return

        stdout, stderr, returncode = result
        self._write_output("")

        if tool == "nmap":
            if stdout:
                self._write_output(
                    "Parsed nmap result:\n" + self._parse_nmap_output(stdout)
                )
            else:
                self._write_output("No stdout returned by nmap.")
        else:
            if stdout:
                self._write_output(f"{tool} output:\n" + stdout.strip())
            else:
                self._write_output(f"No stdout returned by {tool}.")

        if stderr:
            self._write_output("stderr:\n" + stderr.strip())

        if returncode != 0:
            self._write_output(f"{tool} exited with status code {returncode}.")

        self.current_tool = None

    def _normalize_tool_command(self, tool: str, command_text: str) -> list[str]:
        parts = shlex.split(command_text)
        if not parts:
            raise ValueError("empty command")

        if parts[0] != tool:
            parts.insert(0, tool)

        return parts

    async def _run_command(self, tool: str, command: list[str]):
        try:
            completed = await asyncio.to_thread(
                subprocess.run,
                command,
                capture_output=True,
                text=True,
                check=False,
                timeout=180,
            )
        except FileNotFoundError:
            self._write_output(f"{tool} not found. Install {tool} first.")
            return None
        except subprocess.TimeoutExpired:
            self._write_output(f"{tool} command timed out after 180 seconds.")
            return None
        except OSError as exc:
            self._write_output(f"Failed to run {tool}: {exc}")
            return None

        return completed.stdout, completed.stderr, completed.returncode

    def _parse_nmap_output(self, raw_output: str) -> str:
        host = "unknown"
        host_status = ""
        done_line = ""
        ports: list[tuple[str, str, str, str]] = []

        for line in raw_output.splitlines():
            stripped = line.strip()

            host_match = re.match(r"^Nmap scan report for\s+(.+)$", stripped)
            if host_match:
                host = host_match.group(1)
                continue

            if stripped.startswith("Host is "):
                host_status = stripped
                continue

            port_match = re.match(
                r"^(\d+/(?:tcp|udp|sctp))\s+(\S+)\s+(\S+)(?:\s+(.*))?$",
                stripped,
            )
            if port_match:
                port, state, service, version = port_match.groups()
                ports.append((port, state, service, version or ""))
                continue

            if stripped.startswith("Nmap done:"):
                done_line = stripped

        lines = [f"Target: {host}"]
        if host_status:
            lines.append(f"Status: {host_status}")

        if ports:
            lines.append("Open/visible ports:")
            for port, state, service, version in ports:
                detail = f"- {port}: {state} {service}"
                if version:
                    detail += f" ({version})"
                lines.append(detail)
        else:
            lines.append("No port rows detected in output.")

        if done_line:
            lines.append(done_line)

        return "\n".join(lines)

    def _write_output(self, text: str) -> None:
        output = self.query_one("#output", RichLog)
        output.write(text)


if __name__ == "__main__":
    app = COMP7904()
    app.run()