scanner_cli.py

import typer, asyncio, ipaddress, sys
from typing import Literal
from datetime import datetime

from rich.progress import track, Progress, MofNCompleteColumn, TextColumn, BarColumn, SpinnerColumn
from rich.table import Table, Column
from rich.console import Console, Group
from rich.live import Live
from rich.panel import Panel

from scapy.all import get_if_list, get_working_if, resolve_iface

from service.core import ScannerEngine
from service.models import HostDiscoveryResult
from service.utils import get_netmask_cidr

app = typer.Typer(rich_markup_mode="markdown")
console = Console()

class AsyncKeyReader:
    def __init__(self):
        self.os_name = sys.platform
        if self.os_name == 'win32':
            import msvcrt
            self.msvcrt = msvcrt
        else:
            import termios, tty
            self.termios = termios
            self.tty = tty

    def __enter__(self):
        if self.os_name != 'win32':
            self.fd = sys.stdin.fileno()
            self.old_settings = self.termios.tcgetattr(self.fd)
            self.tty.setcbreak(self.fd)
        return self

    def __exit__(self, type, value, traceback):
        if self.os_name != 'win32':
            self.termios.tcsetattr(self.fd, self.termios.TCSADRAIN, self.old_settings)

    def get_key(self) -> str | None:
        """非阻塞读取单个字符"""
        if self.os_name == 'win32':
            if self.msvcrt.kbhit():
                ch = self.msvcrt.getch()
                try:
                    return ch.decode('utf-8').lower()
                except:
                    return None
        else:
            import select
            dr, dw, de = select.select([sys.stdin], [], [], 0)
            if dr:
                return sys.stdin.read(1).lower()
        return None

@app.command()
def list_ifaces():
    '''列出所有的网络接口'''
    table = Table(
        Column("Name", overflow='fold'),
        Column("Desc", overflow='fold'),
        # Column("NetName", overflow='fold'),
        "MAC",
        Column("IP", overflow='fold'),
        title="NetworkInterfaces",
        show_lines=True
    )
    working_if = get_working_if()
    for iface_name in get_if_list():
        iface = resolve_iface(iface_name)
        table.add_row(
            iface.name,
            iface.description,
            # iface.network_name,
            iface.mac,
            '\n'.join(('\n'.join(iface.ips[4]), '\n'.join(iface.ips[6]))),
            style="bold yellow" if iface.network_name == working_if.network_name else ""
        )
    console.print(table)

@app.command()
def scan_ports(
    target_ip: str = typer.Argument(help="目标IP"),
    ports: str = typer.Argument(help='要扫描的端口，格式为"21, 22, 80"'),
    iface: str = typer.Option('$default', help="使用的接口的名称"),
    workers: int = typer.Option(50, min=1, help='最大并发线程数'),
    timeout: float = typer.Option(3.0, min=0.5, help='超时时间'),
    proto: Literal['tcp', 'udp'] = typer.Option('tcp', help='扫描所用协议'),
    method: Literal['syn', 'fin', 'xmas'] = typer.Option('syn', help='扫描方式')
):
    """通过网络接口对目标IP进行TCP端口扫描"""
    interface_name = iface if iface != '$default' else get_working_if().name
    iports = [int(p.strip()) for p in ports.split(',')]

    found_results: list[tuple[int, str, str]] = []

    is_finished = False
    def generate_renderable(prog, is_finished: bool = False):
        """生成 Live 显示的组件组合"""
        table = Table(title=f"Scanning [b green]{target_ip}[/]", expand=False, show_lines=True)
        table.add_column("PROTOCOL", style="magenta", justify="center")
        table.add_column("PORT", style="cyan", justify="right")
        table.add_column("STATE", style="bold green")
        #table.add_column("SERVICE", style="dim")

        for res in found_results:
            table.add_row(res['proto'], str(res['port']), res['state'])

        return Group(
            table,
            Panel(prog, title="Overall Progress", border_style="dim", expand=False)
        ) if not is_finished else table

    async def _scan_ports_internal():
        scanner = ScannerEngine(interface_name, workers)
        
        progress = Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            MofNCompleteColumn(),
            SpinnerColumn(),
        )

        scan_task = progress.add_task("[bold yellow]Scanning", total=len(iports))

        with Live(generate_renderable(progress), refresh_per_second=10) as live:
            results_set = set()
            
            async for portResult in scanner.scan_ports(target_ip, iports, timeout, proto, method):
                found_results.append({
                    "port": portResult.port,
                    "proto": portResult.protocol.upper(),
                    "state": portResult.state.capitalize()
                })
                results_set.add(portResult.port)
                
                progress.advance(scan_task, 1)
                live.update(generate_renderable(progress))

            
            for port in iports:
                if port not in results_set:
                    if proto == 'tcp':
                        found_results.append({
                            "port": port,
                            "proto": "TCP",
                            "state": "Open|Filtered" if method != 'syn' else "Unknown"
                        })
                    elif proto == 'udp':
                        found_results.append({
                            "port": port,
                            "proto": "UDP",
                            "state": "Open|Filtered"
                        })
                    progress.advance(scan_task, 1)
                    live.update(generate_renderable(progress))
        
            live.update(generate_renderable(progress, True))

    asyncio.run(_scan_ports_internal())

@app.command()
def broadcast(
    ip_version: Literal['4', '6', '46'] = typer.Argument(help="要广播的IP协议版本"),
    interface: str = typer.Option('$default', help="使用的接口的名称"),
    workers: int = typer.Option(50, min=1, help='最大并发线程数'),
    timeout: float = typer.Option(10.0, min=1.0, help='响应监视时间'),
    manuf: bool = typer.Option(False, help='显示检测的制造商'),
    manuf_update: bool = typer.Option(False, help='更新制造商信息'),
):
    '''在局域网广播实现设备发现'''
    interface_name = interface if interface != '$default' else get_working_if().name

    found_results: dict[str, HostDiscoveryResult] = {}

    def generate_renderable(prog, is_finished: bool = False, scanner: ScannerEngine | None = None):
        """生成 Live 显示的组件组合"""
        table = Table(title=f"Broadcast on [b yellow]{interface_name}[/]", expand=False, show_lines=True)
        table.add_column("IP", style="bold green")
        table.add_column("MAC", style="magenta")
        table.add_column("METHODS", style="cyan")
        #table.add_column("SERVICE", style="dim")
        if manuf and scanner and scanner.manuf2:
            table.add_column("MANUF")

        for res in found_results.values():
            if manuf and scanner and scanner.manuf2:
                table.add_row(res.ip, res.mac, '\n'.join(res.methods), res.manuf or 'Unknown')
            else:
                table.add_row(res.ip, res.mac, '\n'.join(res.methods))

        return Group(
            table,
            Panel(prog, title="Overall Progress", border_style="dim", expand=False)
        ) if not is_finished else table

    async def _broadcast_internal():
        scanner = ScannerEngine(interface_name, workers, manuf_init=manuf, manuf_update=manuf_update)
        
        progress = Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
        )

        scan_task = progress.add_task("[bold yellow]Listening", total=None)

        with Live(generate_renderable(progress), refresh_per_second=10) as live:            
            async for hostResult in scanner.discover_network('4' in ip_version, '6' in ip_version, timeout):
                if hostResult.ip in found_results:
                    for method in hostResult.methods:
                        found_results[hostResult.ip].methods.add(method)
                else:
                    found_results[hostResult.ip] = hostResult
                
                progress.advance(scan_task, 1)
                live.update(generate_renderable(progress, scanner=scanner))
            
            live.update(generate_renderable(progress, True, scanner))

    asyncio.run(_broadcast_internal())

@app.command()
def scan_cidr(
    cidr: str = typer.Argument('', help="要扫描的地址段"),
    interface: str = typer.Option('$default', help="使用的接口的名称"),
    workers: int = typer.Option(50, min=1, help='最大并发线程数'),
    timeout: float = typer.Option(10.0, min=1.0, help='响应监视时间'),
    batch_size: int = typer.Option(100, min=1, help='每批次发送包数'),
    inter: float = typer.Option(0.002, min=0.0, help='包发送间隔'),
    methods: str = typer.Option('arp', help='扫描方式，可选*arp* *icmp*，也可同时选择多个，如"*arp,icmp*"'),
    manuf: bool = typer.Option(False, help='显示检测的制造商'),
    manuf_update: bool = typer.Option(False, help='更新制造商信息'),
):
    '''通过全量扫描实现设备发现(仅IPv4)'''
    interface_name = interface if interface != '$default' else get_working_if().name

    if not cidr:
        iface = resolve_iface(interface_name)
        cidr = get_netmask_cidr(iface)
        if len(cidr) == 0:
            raise ValueError('Cannot find cidr')
        else:
            cidr = cidr[0]

    test_net = ipaddress.ip_network(cidr, strict=False)
    if not isinstance(test_net, ipaddress.IPv4Network):
        raise ValueError('CIDR SCAN only support IPv4')

    methods = set(method.strip().lower() for method in methods.split(',') if method.strip().lower() in ('arp', 'icmp'))
    if not methods:
        raise ValueError('No enough supported methods')

    found_results: dict[str, HostDiscoveryResult] = {}

    def generate_renderable(prog, is_finished: bool = False, scanner: ScannerEngine | None = None):
        """生成 Live 显示的组件组合"""
        table = Table(title=f"CIDR-Scan on [b yellow]{interface_name}[/], {cidr}", expand=False, show_lines=True)
        table.add_column("IP", style="bold green")
        table.add_column("MAC", style="magenta")
        table.add_column("METHODS", style="cyan")
        #table.add_column("SERVICE", style="dim")
        if manuf and scanner and scanner.manuf2:
            table.add_column("MANUF")

        for res in found_results.values():
            if manuf and scanner and scanner.manuf2:
                table.add_row(res.ip, res.mac, '\n'.join(res.methods), res.manuf or 'Unknown')
            else:
                table.add_row(res.ip, res.mac, '\n'.join(res.methods))

        return Group(
            table,
            Panel(prog, title="Overall Progress", border_style="dim", expand=False)
        ) if not is_finished else table

    async def _scan_cidr_internal():
        scanner = ScannerEngine(interface_name, workers, manuf_init=manuf, manuf_update=manuf_update)
        
        progress = Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
        )

        scan_task = progress.add_task("[bold yellow]Listening", total=None)

        with Live(generate_renderable(progress), refresh_per_second=10) as live:            
            async for hostResult in scanner.scan_network_cidr(cidr, timeout, batch_size, inter, methods):
                if hostResult.ip in found_results:
                    for method in hostResult.methods:
                        found_results[hostResult.ip].methods.add(method)
                else:
                    found_results[hostResult.ip] = hostResult
                
                progress.advance(scan_task, 1)
                live.update(generate_renderable(progress, scanner=scanner))
            
            live.update(generate_renderable(progress, True, scanner))

    asyncio.run(_scan_cidr_internal())

@app.command()
def monitor(
    cidr: str = typer.Option(None, help="主动全量扫描时使用的默认CIDR (IPv4)"),
    interface: str = typer.Option('$default', help="使用的接口的名称"),
    workers: int = typer.Option(50, min = 1, help='最大并发线程数'),
    manuf_instead: bool = typer.Option(False, help='显示检测的制造商而不是响应的主机名'),
    manuf_update: bool = typer.Option(False, help='更新制造商信息'),
    alive_threshold: int = typer.Option(30, min=10, help='多久未响应判定设备离线 (秒)'),
    alive_inter: float = typer.Option(10.0, min=1.0, help='定期确认设备是否依然在线 (秒)'),
    cidr_timeout: float = typer.Option(10.0, min=1, help='主动全量扫描的超时时间 (秒)'),
    cidr_batch: int = typer.Option(100, min=1, help='每批次发送包数'),
    cidr_inter: float = typer.Option(0.002, min=0.0, help='包发送间隔 (秒)'),
    cidr_methods: str = typer.Option('arp', help='扫描方式，可选*arp* *icmp*，也可同时选择多个，如"*arp,icmp*"'),
    broadcast_timeout: float = typer.Option(10.0, min=1, help='主动广播的超时时间 (秒)')
):
    '''
    [交互式] 动态网络监听器
    
    功能：
    1. 持续被动监听网络中的广播/多播流量。
    2. 定期对已发现设备进行心跳检测 (Active/Offline)。
    3. 支持按键触发主动扫描。
    
    控制按键：
    * [b] - 触发局域网广播发现
    * [s] - 触发CIDR网段扫描 (需指定 --cidr 或自动推断)
    * [c] - 清除离线设备
    * [q] - 退出
    '''
    interface_name = interface if interface != '$default' else get_working_if().name
    interface_mac = resolve_iface(interface_name).mac
    
    # 状态存储
    # key: ip, value: HostDiscoveryResult
    found_hosts: dict[str, HostDiscoveryResult] = {}
    # key: ip, value: (is_online: bool, last_seen: datetime)
    host_status: dict[str, tuple[bool, datetime]] = {}
    # 日志消息列表
    logs: list[str] = []

    # 用于传递给心跳检测的列表 (动态更新)
    heartbeat_targets: list[HostDiscoveryResult] = []

    def add_log(msg: str):
        time_str = datetime.now().strftime("%H:%M:%S")
        logs.append(f"[{time_str}] {msg}")
        if len(logs) > 8: # 保持日志简短
            logs.pop(0)

    def update_host(result: HostDiscoveryResult, source: str = "Passive"):
        """更新主机信息"""
        is_new = result.ip not in found_hosts
        
        if is_new:
            found_hosts[result.ip] = result
            host_status[result.ip] = (True, datetime.now())
            heartbeat_targets.append(result) # 加入心跳监控名单
            add_log(f"[bold green]New Device[/]: {result.ip} ({source})")
        else:
            existing = found_hosts[result.ip]
            if result.hostname and not existing.hostname:
                existing.hostname = result.hostname
                add_log(f"[bold cyan]Info Update[/]: {result.ip} hostname -> {result.hostname}")
            if result.mac and not existing.mac:
                existing.mac = result.mac
            for m in result.methods:
                existing.methods.add(m)
            
            host_status[result.ip] = (True, datetime.now())

    def generate_dashboard():
        """生成主界面"""
        table = Table(expand=True, title=f"Monitor on [bold yellow]{interface_name}[/]", box=None)
        table.add_column("IP Address", style="bold")
        table.add_column("MAC Address", style="dim")
        table.add_column("Hostname" if not manuf_instead else "Manuf")
        table.add_column("Methods", style="magenta")
        table.add_column("Status", justify="right")
        table.add_column("Last Seen", style="dim", justify="right")

        def get_sort_key(ip_str):
            try:
                ip_obj = ipaddress.ip_address(ip_str)
                is_online = host_status[ip_str][0]
                
                return (
                    not is_online,
                    ip_obj.version,
                    ip_obj 
                )
            except ValueError:
                return (True, 999, 0)

        sorted_ips = sorted(
            found_hosts.keys(), 
            key=get_sort_key
        )

        online_count = 0
        for ip in sorted_ips:
            host = found_hosts[ip]
            is_online, last_time = host_status[ip]
            if is_online: online_count += 1
            
            status_style = "bold green" if is_online else "red"
            status_text = "● Online" if is_online else "○ Offline"
            
            time_delta = (datetime.now() - last_time).seconds
            time_str = f"{time_delta}s ago" if time_delta < 60 else f"{time_delta//60}m ago"

            is_host = interface_mac and host.mac and interface_mac == host.mac
            table.add_row(
                host.ip if not is_host else f"[yellow]{host.ip}[/yellow]",
                host.mac or "-",
                (host.hostname or "") if not manuf_instead else (host.manuf or 'Unknown'),
                ", ".join(list(host.methods)[:3]), # 只显示前3个方法避免过长
                f"[{status_style}]{status_text}[/]",
                time_str,
            )

        stats_panel = Panel(
            f"Total: {len(found_hosts)} | Online: [green]{online_count}[/] | Offline: [red]{len(found_hosts)-online_count}[/]",
            title="Statistics",
            border_style="blue"
        )

        log_content = "\n".join(logs) if logs else "[dim]Waiting for events...[/]"
        log_panel = Panel(log_content, title="Event Log", border_style="yellow", height=10)

        help_text = "[b]Controls[/]: [bold magenta]b[/]roadcast discovery | [bold magenta]s[/]can cidr | [bold magenta]c[/]lear offline | [bold red]q[/]uit"
        
        return Group(stats_panel, table, log_panel, help_text)

    async def _monitor_loop(manuf_instead: bool, manuf_update: bool = False):
        engine = ScannerEngine(interface_name, workers, manuf_init=manuf_instead, manuf_update=manuf_update)
        if manuf_instead and not engine.manuf2:
            manuf_instead = False
        add_log(f"Engine started. Listening for traffic...")

        loop = asyncio.get_running_loop()
        background_tasks: set[asyncio.Task] = set()


        async def task_passive():
            """被动监听协程"""
            async for res in engine.start_passive_monitor():
                update_host(res, "Passive")

        async def task_heartbeat():
            """心跳保活协程"""
            while True:
                try:
                    async for ip, is_alive in engine.monitor_heartbeat(heartbeat_targets, interval=alive_inter):
                        if ip in host_status:
                            current_status, last_seen = host_status[ip]
                            if is_alive:
                                host_status[ip] = (True, datetime.now())
                            else:
                                if (datetime.now() - last_seen).seconds > alive_threshold:
                                    if current_status: # 状态变更
                                        host_status[ip] = (False, last_seen)
                                        add_log(f"[red]Device Offline[/]: {ip}")
                except asyncio.CancelledError:
                    break
                except:
                    pass

        async def task_active_broadcast():
            """一次性广播任务"""
            add_log("[yellow]Starting Active Broadcast...[/]")
            async for res in engine.discover_network(scan_v4=True, scan_v6=True, timeout=broadcast_timeout):
                update_host(res, "Broadcast")
            add_log("[green]Broadcast Finished.[/]")

        async def task_active_cidr():
            """一次性CIDR扫描任务"""
            target_cidr = cidr
            if not target_cidr:
                # 自动推断 CIDR
                from service.utils import get_netmask_cidr
                nets = get_netmask_cidr(engine.iface, 4)
                if nets:
                    target_cidr = nets[0]
            
            if not target_cidr:
                add_log("[red]Error: No CIDR specified or found.[/]")
                return

            add_log(f"[yellow]Scanning CIDR: {target_cidr}...[/]")
            async for res in engine.scan_network_cidr(target_cidr, timeout=cidr_timeout, batch_size=cidr_batch, inter=cidr_inter, methods=set(m.strip().lower() for m in cidr_methods.split(','))):
                update_host(res, "CIDR-Scan")
            add_log("[green]CIDR Scan Finished.[/]")

        t_pass = asyncio.create_task(task_passive())
        t_heart = asyncio.create_task(task_heartbeat())
        background_tasks.add(t_pass)
        background_tasks.add(t_heart)

        with Live(generate_dashboard(), refresh_per_second=4, screen=True) as live:
            with AsyncKeyReader() as key_reader:
                t = None
                while True:
                    live.update(generate_dashboard())
                    
                    key = key_reader.get_key()
                    if key == 'q':
                        break
                    elif key == 'b':
                        if not t or t.done():
                            t = asyncio.create_task(task_active_broadcast())
                            background_tasks.add(t)
                            t.add_done_callback(background_tasks.discard)
                        else:
                             add_log("[red]Last task is still running[/]")
                    elif key == 's':
                        if not t or t.done():
                            t = asyncio.create_task(task_active_cidr())
                            background_tasks.add(t)
                            t.add_done_callback(background_tasks.discard)
                        else:
                             add_log("[red]Last task is still running[/]")
                    elif key == 'c':
                        to_remove = [ip for ip, stat in host_status.items() if not stat[0]]
                        for ip in to_remove:
                            del found_hosts[ip]
                            del host_status[ip]
                            # 从心跳名单移除
                            for i in range(len(heartbeat_targets)-1, -1, -1):
                                if heartbeat_targets[i].ip == ip:
                                    heartbeat_targets.pop(i)
                        add_log(f"Cleared {len(to_remove)} offline devices.")

                    await asyncio.sleep(0.1)

        # 清理任务
        for t in background_tasks:
            t.cancel()
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            transient=True
        ) as progress:
            progress.add_task('Please wait for tasks to exit')
            while any(not t.done() for t in background_tasks):
                await asyncio.sleep(0.1)
    try:
        asyncio.run(_monitor_loop(manuf_instead, manuf_update))
    except KeyboardInterrupt:
        pass

if __name__ == "__main__":
    app()