portscout.py

#!/usr/bin/env python3
"""
PortScout TUI - interactive terminal UI for port scanning & recommendations.

Features:
 - Scan a port range for free/used ports (attempts to bind locally)
 - Shows used ports (with PID/process name if psutil installed)
 - Shows free ports (compact, with contiguous grouping)
 - Re-scan (r), change scan range (g), exclude ports (x), toggle contiguous preference (c)
 - Reserve a selected port (press 'v'): reservations saved to ~/.portscout/reservations.yaml
 - Quit with 'q' or Ctrl-C
 - Responsive, non-blocking scanning using ThreadPoolExecutor + asyncio

Dependencies: textual, psutil (optional but recommended), pyyaml
Install: pip install textual psutil pyyaml

Fixes applied vs original:
  1. Input dialogs now use proper Textual ModalScreen push/pop instead of
     calling action_submit() immediately (which never waited for the user).
  2. scan_ports_async now uses asyncio.gather() so port→result pairs stay
     correctly ordered (as_completed returns in *completion* order, not
     submission order, causing the original zip() to mis-label ports).
  3. action_quit uses self.exit() instead of the deprecated self.shutdown().
  4. Added basic CSS for a readable two-column layout.
"""

import asyncio
import socket
import concurrent.futures
import functools
import yaml
from pathlib import Path
from typing import List, Set, Tuple

from textual.app import App, ComposeResult
from textual.widgets import Header, Footer, Static, Input, Label, Button
from textual.containers import Horizontal, Vertical
from textual.reactive import reactive
from textual.screen import ModalScreen
from textual import events

# Optional psutil
try:
    import psutil
except Exception:
    psutil = None

RESERVE_DIR = Path.home() / ".portscout"
RESERVE_FILE = RESERVE_DIR / "reservations.yaml"

# ──────────────────────────────────────────────
# Low-level scanning helpers
# ──────────────────────────────────────────────

def try_bind(port: int, host: str = "0.0.0.0", timeout: float = 0.35) -> bool:
    """Return True if we can bind to *port* (i.e. it is free)."""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(timeout)
    try:
        s.bind((host, port))
        s.listen(1)
        s.close()
        return True
    except Exception:
        return False


def get_used_ports_with_info() -> dict:
    """Return {port: (pid, pname)} for all listening sockets via psutil."""
    used: dict = {}
    if not psutil:
        return used
    for conn in psutil.net_connections(kind="inet"):
        if not conn.laddr:
            continue
        port = conn.laddr.port
        if port == 0:
            continue
        if port not in used:
            pid = conn.pid
            pname = None
            try:
                p = psutil.Process(pid) if pid else None
                pname = p.name() if p else None
            except Exception:
                pname = None
            used[port] = (pid, pname)
    return used


def contiguous_ranges(lst: List[int]) -> List[Tuple[int, int]]:
    """Collapse a sorted list of ints into (start, end) range tuples."""
    if not lst:
        return []
    lst = sorted(lst)
    ranges: List[Tuple[int, int]] = []
    start = prev = lst[0]
    for x in lst[1:]:
        if x == prev + 1:
            prev = x
            continue
        ranges.append((start, prev))
        start = prev = x
    ranges.append((start, prev))
    return ranges


def parse_exclude(s: str) -> Set[int]:
    """Parse a string like '22,80,8000-8010' into a set of ints."""
    out: Set[int] = set()
    if not s:
        return out
    for part in (p.strip() for p in s.split(",") if p.strip()):
        if "-" in part:
            a, b = part.split("-", 1)
            out.update(range(int(a), int(b) + 1))
        else:
            out.add(int(part))
    return out


# ──────────────────────────────────────────────
# Async scanning
# ──────────────────────────────────────────────

async def scan_ports_async(
    start: int,
    end: int,
    threads: int = 300,
    exclude: Set[int] = None,
    host: str = "0.0.0.0",
) -> Tuple[List[int], List[int]]:
    """
    Scan ports in [start, end] (excluding *exclude*) and return
    (free_ports, used_ports).

    FIX: original code used asyncio.as_completed() zipped with the ports list.
    as_completed() yields futures in *completion* order, not submission order,
    so the zip() would mis-label which result belonged to which port.
    We now use asyncio.gather() which preserves submission order.
    """
    exclude = exclude or set()
    ports = [p for p in range(start, end + 1) if p not in exclude]
    free: List[int] = []
    used: List[int] = []

    loop = asyncio.get_event_loop()
    with concurrent.futures.ThreadPoolExecutor(
        max_workers=min(threads, len(ports) or 1)
    ) as executor:
        tasks = [
            loop.run_in_executor(executor, functools.partial(try_bind, p, host))
            for p in ports
        ]
        # gather() keeps results in the same order as tasks/ports
        results = await asyncio.gather(*tasks)

    for port, is_free in zip(ports, results):
        (free if is_free else used).append(port)

    free.sort()
    used.sort()
    return free, used


# ──────────────────────────────────────────────
# Reservation helpers
# ──────────────────────────────────────────────

def load_reservations() -> dict:
    try:
        if RESERVE_FILE.exists():
            return yaml.safe_load(RESERVE_FILE.read_text()) or {}
    except Exception:
        pass
    return {}


def save_reservation(port: int, note: str = "") -> None:
    RESERVE_DIR.mkdir(parents=True, exist_ok=True)
    data = load_reservations()
    data[str(port)] = {"port": port, "note": note}
    RESERVE_FILE.write_text(yaml.safe_dump(data))


# ──────────────────────────────────────────────
# Modal screens  (FIX: replace broken action_submit() pattern)
# ──────────────────────────────────────────────

class SingleInputModal(ModalScreen[str | None]):
    """
    Generic single-field modal.  Resolves with the entered string on
    Enter/Submit, or None if the user presses Escape.
    """

    DEFAULT_CSS = """
    SingleInputModal {
        align: center middle;
    }
    SingleInputModal > Vertical {
        width: 60;
        height: auto;
        background: $surface;
        border: thick $primary;
        padding: 1 2;
    }
    SingleInputModal Label {
        margin-bottom: 1;
    }
    SingleInputModal #modal-buttons {
        margin-top: 1;
        height: auto;
    }
    SingleInputModal Button {
        margin-right: 1;
    }
    """

    def __init__(self, prompt: str, default: str = "", **kwargs):
        super().__init__(**kwargs)
        self._prompt = prompt
        self._default = default

    def compose(self) -> ComposeResult:
        with Vertical():
            yield Label(self._prompt)
            yield Input(value=self._default, id="modal-input")
            with Horizontal(id="modal-buttons"):
                yield Button("OK", variant="primary", id="ok")
                yield Button("Cancel", id="cancel")

    def on_mount(self) -> None:
        self.query_one("#modal-input", Input).focus()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "ok":
            self.dismiss(self.query_one("#modal-input", Input).value)
        else:
            self.dismiss(None)

    def on_input_submitted(self, event: Input.Submitted) -> None:
        self.dismiss(event.value)

    def on_key(self, event: events.Key) -> None:
        if event.key == "escape":
            self.dismiss(None)


class RangeModal(ModalScreen[Tuple[int, int] | None]):
    """Two-field modal for changing the scan range."""

    DEFAULT_CSS = """
    RangeModal {
        align: center middle;
    }
    RangeModal > Vertical {
        width: 60;
        height: auto;
        background: $surface;
        border: thick $primary;
        padding: 1 2;
    }
    RangeModal Label {
        margin-bottom: 0;
    }
    RangeModal Input {
        margin-bottom: 1;
    }
    RangeModal #modal-buttons {
        margin-top: 1;
        height: auto;
    }
    RangeModal Button {
        margin-right: 1;
    }
    """

    def __init__(self, start: int, end: int, **kwargs):
        super().__init__(**kwargs)
        self._start = start
        self._end = end

    def compose(self) -> ComposeResult:
        with Vertical():
            yield Label("Start port:")
            yield Input(value=str(self._start), id="start-input")
            yield Label("End port:")
            yield Input(value=str(self._end), id="end-input")
            with Horizontal(id="modal-buttons"):
                yield Button("OK", variant="primary", id="ok")
                yield Button("Cancel", id="cancel")

    def on_mount(self) -> None:
        self.query_one("#start-input", Input).focus()

    def _parse_result(self) -> Tuple[int, int] | None:
        sv = self.query_one("#start-input", Input).value
        ev = self.query_one("#end-input", Input).value
        if sv.isdigit() and ev.isdigit():
            return int(sv), int(ev)
        return None

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "ok":
            self.dismiss(self._parse_result())
        else:
            self.dismiss(None)

    def on_key(self, event: events.Key) -> None:
        if event.key == "escape":
            self.dismiss(None)


# ──────────────────────────────────────────────
# Main Textual app
# ──────────────────────────────────────────────

APP_CSS = """
Screen {
    layout: vertical;
}

#columns {
    layout: horizontal;
    height: 1fr;
}

#left-col {
    width: 1fr;
    border: solid $primary-darken-2;
    padding: 0 1;
    overflow-y: auto;
}

#right-col {
    width: 1fr;
    border: solid $primary-darken-2;
    padding: 0 1;
    overflow-y: auto;
}

#used-info {
    height: auto;
    border-top: dashed $primary-darken-3;
    padding: 0 1;
    color: $text-muted;
}

#recommend-box {
    height: auto;
    border-top: dashed $primary-darken-3;
    padding: 0 1;
    color: $success;
}
"""


class InfoPanel(Static):
    def update_text(self, text: str) -> None:
        self.update(text)


class PortList(Static):
    def __init__(self, title: str, **kwargs):
        super().__init__(**kwargs)
        self._title = title

    def update_list(self, items: List[int]) -> None:
        lines = [f"[b]{self._title}[/b] ({len(items)} ports)\n"]
        display = items[:500]
        for a, b in contiguous_ranges(display):
            lines.append(f"  {a}" if a == b else f"  {a}–{b}")
        if len(items) > 500:
            lines.append(f"\n  … and {len(items) - 500} more")
        self.update("\n".join(lines))


class PortScoutApp(App):
    CSS = APP_CSS
    TITLE = "PortScout"
    SUB_TITLE = "terminal port scanner"

    BINDINGS = [
        ("r", "rescan", "Rescan"),
        ("g", "change_range", "Change range"),
        ("x", "set_exclude", "Set exclusions"),
        ("c", "toggle_contig", "Toggle contiguous"),
        ("v", "reserve", "Reserve port"),
        ("h", "show_help", "Help"),
        ("q", "quit", "Quit"),
    ]

    start_port: reactive[int] = reactive(1024)
    end_port: reactive[int] = reactive(9999)
    count: int = 8
    contiguous_pref: reactive[bool] = reactive(False)
    exclude_set: reactive[Set[int]] = reactive(set)
    scanning: reactive[bool] = reactive(False)
    last_free: List[int] = []
    last_used: List[int] = []
    used_info: dict = {}

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        with Horizontal(id="columns"):
            with Vertical(id="left-col"):
                self.used_list = PortList("Used Ports")
                yield self.used_list
                self.used_info_box = InfoPanel("", id="used-info")
                yield self.used_info_box
            with Vertical(id="right-col"):
                self.free_list = PortList("Free Ports")
                yield self.free_list
                self.recommend_box = InfoPanel("", id="recommend-box")
                yield self.recommend_box
        yield Footer()

    async def on_mount(self) -> None:
        self.used_info = get_used_ports_with_info()
        await self.action_rescan()

    # ── scanning ──────────────────────────────

    async def action_rescan(self) -> None:
        if self.scanning:
            return
        self.scanning = True
        self.recommend_box.update(
            f"[yellow]Scanning {self.start_port}–{self.end_port} …[/yellow]"
        )
        try:
            free, used = await scan_ports_async(
                self.start_port,
                self.end_port,
                threads=300,
                exclude=self.exclude_set,
            )
            self.last_free = free
            self.last_used = used
            self.used_info = get_used_ports_with_info()

            self.used_list.update_list(self.last_used)
            lines = ["[b]Used ports – process info (first 50):[/b]"]
            for p in self.last_used[:50]:
                pid, pname = self.used_info.get(p, (None, None))
                lines.append(f"  {p:<6} PID={pid or '-':<7} {pname or '?'}")
            if not self.last_used:
                lines = ["[dim]No used ports found in range.[/dim]"]
            self.used_info_box.update("\n".join(lines))

            self.free_list.update_list(self.last_free)
            self._refresh_recommendations()
        finally:
            self.scanning = False

    def _refresh_recommendations(self) -> None:
        rec = self._recommend(self.last_free, self.count, self.contiguous_pref)
        exc_note = (
            f"  [dim]({len(self.exclude_set)} ports excluded)[/dim]"
            if self.exclude_set
            else ""
        )
        contig_note = "  [dim](contiguous mode)[/dim]" if self.contiguous_pref else ""
        if rec:
            ports_str = ", ".join(map(str, rec))
            warn = (
                "\n  [red]⚠ Some recommended ports are <1024 (privileged).[/red]"
                if any(p < 1024 for p in rec)
                else ""
            )
            self.recommend_box.update(
                f"[b]Recommended:[/b] {ports_str}{exc_note}{contig_note}{warn}"
            )
        else:
            self.recommend_box.update(
                f"[dim]No free ports found in range {self.start_port}–{self.end_port}.[/dim]"
                f"{exc_note}{contig_note}"
            )

    @staticmethod
    def _recommend(free: List[int], count: int, contiguous: bool) -> List[int]:
        if not free:
            return []
        if contiguous:
            for a, b in contiguous_ranges(free):
                if (b - a + 1) >= count:
                    return list(range(a, a + count))
        return free[:count]

    # ── actions (use modal screens) ───────────

    async def action_change_range(self) -> None:
        """FIX: push a proper ModalScreen; await its result via callback."""

        def handle_result(result: Tuple[int, int] | None) -> None:
            if result is None:
                return
            new_start, new_end = result
            new_start = max(1, new_start)
            new_end = min(65535, new_end)
            if new_start > new_end:
                self.recommend_box.update("[red]Invalid range (start > end). Unchanged.[/red]")
                return
            self.start_port = new_start
            self.end_port = new_end
            self.app.call_later(self.action_rescan)

        await self.push_screen(RangeModal(self.start_port, self.end_port), handle_result)

    async def action_set_exclude(self) -> None:
        """FIX: push a proper ModalScreen for exclude input."""

        def handle_result(value: str | None) -> None:
            if value is None:
                return
            value = value.strip()
            if not value:
                self.exclude_set = set()
                self.recommend_box.update("[dim]Exclusions cleared.[/dim]")
            else:
                try:
                    self.exclude_set = parse_exclude(value)
                    self.recommend_box.update(
                        f"Excluding {len(self.exclude_set)} ports. Press [b]r[/b] to rescan."
                    )
                except Exception as exc:
                    self.recommend_box.update(f"[red]Error parsing exclusions: {exc}[/red]")

        await self.push_screen(
            SingleInputModal(
                "Exclude ports (e.g. 22,80,8000-8010) — leave blank to clear:",
                default="",
            ),
            handle_result,
        )

    async def action_toggle_contig(self) -> None:
        self.contiguous_pref = not self.contiguous_pref
        state = "[green]ON[/green]" if self.contiguous_pref else "[dim]OFF[/dim]"
        self.recommend_box.update(
            f"Contiguous preference {state}. Press [b]r[/b] to rescan."
        )

    async def action_reserve(self) -> None:
        rec = self._recommend(self.last_free, self.count, self.contiguous_pref)
        if not rec:
            self.recommend_box.update("[red]No recommended free ports to reserve. Rescan first.[/red]")
            return
        port_to_reserve = rec[0]

        def handle_result(note: str | None) -> None:
            if note is None:
                return
            save_reservation(port_to_reserve, note.strip())
            self.recommend_box.update(
                f"[green]✓ Reserved port {port_to_reserve}.[/green]\n"
                f"  Saved to {RESERVE_FILE}"
            )

        await self.push_screen(
            SingleInputModal(
                f"Reserve port {port_to_reserve} — optional note:",
                default="",
            ),
            handle_result,
        )

    async def action_show_help(self) -> None:
        help_text = (
            "[b]Keyboard shortcuts[/b]\n\n"
            "  [b]r[/b]  Rescan the current port range\n"
            "  [b]g[/b]  Change scan range\n"
            "  [b]x[/b]  Set port exclusions\n"
            "  [b]c[/b]  Toggle contiguous-block preference\n"
            "  [b]v[/b]  Reserve the top recommended port\n"
            "  [b]h[/b]  Show this help\n"
            "  [b]q[/b]  Quit"
        )
        self.recommend_box.update(help_text)

    # FIX: use self.exit() instead of deprecated self.shutdown()
    def action_quit(self) -> None:
        self.exit()


if __name__ == "__main__":
    PortScoutApp().run()