"""
FastRDP-NG: Blazing-fast async RDP scanner.
Connects to thousands of IPs concurrently to find live RDP hosts.
"""

import asyncio
import logging
import time
from typing import Set, Tuple, List

logger = logging.getLogger("FastRDP-NG")

# RDP ports to scan
RDP_PORTS = [3389, 3390, 3391]

# First 3 bytes of an RDP Negotiation Response (T.125)
# Actual RDP servers respond with 0x03 (TPKT version 3)
RDP_BANNER_SIG = b'\x03\x00'


async def check_rdp_port(
    ip: str,
    port: int,
    timeout: float = 2.0
) -> Tuple[str, int, bool]:
    """
    Rapid async TCP connect check.
    Returns (ip, port, is_open).
    No banner reading - pure connection speed.
    """
    try:
        _, writer = await asyncio.wait_for(
            asyncio.open_connection(ip, port),
            timeout=timeout
        )
        writer.close()
        await writer.wait_closed()
        return (ip, port, True)
    except (asyncio.TimeoutError, ConnectionRefusedError,
            OSError, ConnectionError):
        return (ip, port, False)


async def check_rdp_with_banner(
    ip: str,
    port: int,
    connect_timeout: float = 2.0,
    banner_timeout: float = 1.0
) -> Tuple[str, int, bool]:
    """
    Connect + read initial bytes to confirm it's actually RDP.
    Slower but more accurate - catches non-RDP services on 3389.
    """
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(ip, port),
            timeout=connect_timeout
        )
        try:
            data = await asyncio.wait_for(
                reader.read(4),
                timeout=banner_timeout
            )
            is_rdp = data[:2] == RDP_BANNER_SIG  # TPKT header
            return (ip, port, is_rdp)
        except (asyncio.TimeoutError, ConnectionError, OSError):
            # Connected but no banner = port open but not RDP
            return (ip, port, False)
        finally:
            writer.close()
            await writer.wait_closed()
    except (asyncio.TimeoutError, ConnectionRefusedError,
            OSError, ConnectionError):
        return (ip, port, False)


async def scan_ips(
    ips: List[str],
    ports: List[int] = None,  # type: ignore
    max_concurrent: int = 5000,
    connect_timeout: float = 2.0,
    banner_check: bool = False,
    progress_callback=None,
    live_callback=None
) -> Set[Tuple[str, int]]:
    """
    Scan a list of IPs across given ports as fast as possible.
    Uses asyncio semaphore to control concurrency.

    - progress_callback(checked, total, live_count, rate): called periodically
    - live_callback(ip, port): called IMMEDIATELY when a live host is found

    Returns set of (ip, port) tuples for live RDP hosts.
    """
    if ports is None:
        ports = RDP_PORTS

    sem = asyncio.Semaphore(max_concurrent)
    live_hosts: Set[Tuple[str, int]] = set()
    checked = 0
    total = len(ips) * len(ports)
    start_time = time.time()

    checker = check_rdp_with_banner if banner_check else check_rdp_port

    async def check_one(ip: str, port: int):
        nonlocal checked
        async with sem:
            _, p, alive = await checker(ip, port, connect_timeout)
            if alive:
                live_hosts.add((ip, p))
                # Notify immediately — spray can start right away
                if live_callback:
                    live_callback(ip, p)
            checked += 1
            if progress_callback and checked % 1000 == 0:
                elapsed = time.time() - start_time
                rate = checked / elapsed if elapsed > 0 else 0
                progress_callback(checked, total, len(live_hosts), rate)

    # Fire all tasks concurrently
    tasks = []
    for ip in ips:
        for port in ports:
            tasks.append(asyncio.create_task(check_one(ip, port)))

    # Run in batches to avoid memory issues with massive task lists
    batch_size = 50000
    for i in range(0, len(tasks), batch_size):
        batch = tasks[i:i + batch_size]
        if batch:
            await asyncio.gather(*batch, return_exceptions=True)

    elapsed = time.time() - start_time
    rate = checked / elapsed if elapsed > 0 else 0
    logger.info(
        f"Scan complete: {len(live_hosts)} live hosts from {checked} checks "
        f"in {elapsed:.1f}s ({rate:.0f} checks/sec)"
    )

    return live_hosts