#!/usr/bin/env python3
"""
Port Scanner - 多线程端口扫描工具

支持:
- TCP Connect 扫描
- SYN 扫描 (需要 root)
- 服务指纹识别
- 多线程扫描
- 自定义端口范围

Usage:
    python3 port-scanner.py -H 192.168.1.1 -p 1-1000
    python3 port-scanner.py -H 192.168.1.1 -p 80,443,8080
    python3 port-scanner.py -H 192.168.1.1 --top-ports 100
"""

import argparse
import socket
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Tuple, Optional
import sys


class Colors:
    RED = "\033[91m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    BLUE = "\033[94m"
    CYAN = "\033[96m"
    END = "\033[0m"
    BOLD = "\033[1m"


class PortScanner:
    def __init__(self, threads: int = 100, timeout: float = 1.0):
        self.threads = threads
        self.timeout = timeout
        self.open_ports = []
        self.lock = threading.Lock()

        self.service_banners = {
            21: "FTP",
            22: "SSH",
            23: "Telnet",
            25: "SMTP",
            53: "DNS",
            80: "HTTP",
            110: "POP3",
            135: "RPC",
            139: "NetBIOS",
            143: "IMAP",
            443: "HTTPS",
            445: "SMB",
            993: "IMAPS",
            995: "POP3S",
            1433: "MSSQL",
            1521: "Oracle",
            3306: "MySQL",
            3389: "RDP",
            5432: "PostgreSQL",
            5900: "VNC",
            6379: "Redis",
            8080: "HTTP-Proxy",
            8443: "HTTPS-Alt",
            8888: "HTTP-Alt",
            9000: "PHP-FPM",
            9200: "Elasticsearch",
            27017: "MongoDB",
        }

        self.top_ports = [
            21,
            22,
            23,
            25,
            53,
            80,
            110,
            111,
            135,
            139,
            143,
            443,
            445,
            993,
            995,
            1433,
            1434,
            1723,
            3306,
            3389,
            5432,
            5900,
            6379,
            8000,
            8080,
            8443,
            8888,
            9000,
            9090,
            9200,
            27017,
        ]

    def print_result(self, level: str, msg: str):
        colors = {
            "INFO": Colors.BLUE,
            "SUCCESS": Colors.GREEN,
            "WARNING": Colors.YELLOW,
            "ERROR": Colors.RED,
            "OPEN": Colors.GREEN + Colors.BOLD,
        }
        print(f"{colors.get(level, '')}[{level}]{Colors.END} {msg}")

    def parse_ports(self, port_str: str) -> List[int]:
        """解析端口字符串"""
        ports = set()

        for part in port_str.split(","):
            if "-" in part:
                start, end = part.split("-")
                ports.update(range(int(start), int(end) + 1))
            else:
                ports.add(int(part))

        return sorted(ports)

    def scan_port(self, host: str, port: int) -> Tuple[int, str, Optional[str]]:
        """扫描单个端口"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.timeout)

            result = sock.connect_ex((host, port))

            if result == 0:
                service = self.service_banners.get(port, "Unknown")

                banner = None
                try:
                    sock.send(b"HEAD / HTTP/1.0\r\n\r\n")
                    banner = sock.recv(1024).decode("utf-8", errors="ignore").strip()
                except:
                    pass

                sock.close()
                return port, "open", banner or service

            sock.close()

        except Exception:
            pass

        return port, "closed", None

    def scan_host(
        self, host: str, ports: List[int], verbose: bool = True
    ) -> List[Dict]:
        """扫描主机"""
        results = []
        total = len(ports)

        if verbose:
            self.print_result("INFO", f"开始扫描 {host}: {total} 个端口")

        start_time = time.time()

        with ThreadPoolExecutor(max_workers=self.threads) as executor:
            futures = {
                executor.submit(self.scan_port, host, port): port for port in ports
            }

            completed = 0
            for future in as_completed(futures):
                port, status, banner = future.result()
                completed += 1

                if verbose and completed % 100 == 0:
                    print(
                        f"\r{Colors.CYAN}[*]{Colors.END} 进度: {completed}/{total}",
                        end="",
                        flush=True,
                    )

                if status == "open":
                    result = {
                        "port": port,
                        "status": status,
                        "service": self.service_banners.get(port, "Unknown"),
                        "banner": banner,
                    }
                    results.append(result)

                    with self.lock:
                        self.open_ports.append((host, port))

                    if verbose:
                        print(
                            f"\n{Colors.GREEN}[OPEN]{Colors.END} {host}:{port} - {banner[:50] if banner else self.service_banners.get(port, 'Unknown')}"
                        )

        if verbose:
            print()

        elapsed = time.time() - start_time
        if verbose:
            self.print_result(
                "INFO", f"扫描完成: {len(results)} 个开放端口, 耗时 {elapsed:.2f}s"
            )

        return results


def main():
    parser = argparse.ArgumentParser(description="Port Scanner")
    parser.add_argument("-H", "--host", required=True, help="目标主机")
    parser.add_argument(
        "-p", "--ports", default="1-1000", help="端口范围 (例: 1-1000, 80,443,8080)"
    )
    parser.add_argument("--top-ports", type=int, help="扫描最常用的 N 个端口")
    parser.add_argument("-t", "--threads", type=int, default=100, help="线程数")
    parser.add_argument("--timeout", type=float, default=1.0, help="超时时间")
    parser.add_argument("-v", "--verbose", action="store_true", help="详细输出")

    args = parser.parse_args()

    scanner = PortScanner(threads=args.threads, timeout=args.timeout)

    if args.top_ports:
        ports = scanner.top_ports[: args.top_ports]
    else:
        ports = scanner.parse_ports(args.ports)

    print(f"\n{Colors.BOLD}{'=' * 60}{Colors.END}")
    print(f"{Colors.BOLD}Port Scanner{Colors.END}")
    print(f"{Colors.BOLD}{'=' * 60}{Colors.END}\n")

    scanner.print_result("INFO", f"目标: {args.host}")
    scanner.print_result("INFO", f"端口: {len(ports)} 个")
    scanner.print_result("INFO", f"线程: {args.threads}")

    results = scanner.scan_host(args.host, ports, args.verbose)

    print(f"\n{Colors.BOLD}{'=' * 60}{Colors.END}")
    if results:
        scanner.print_result("SUCCESS", f"发现 {len(results)} 个开放端口:")
        print(f"\n{'PORT':<10} {'SERVICE':<15} {'BANNER'}")
        print("-" * 60)
        for r in sorted(results, key=lambda x: x["port"]):
            banner = r["banner"][:40] if r["banner"] else r["service"]
            print(f"{r['port']:<10} {r['service']:<15} {banner}")
    else:
        scanner.print_result("INFO", "未发现开放端口")
    print(f"{Colors.BOLD}{'=' * 60}{Colors.END}\n")


if __name__ == "__main__":
    main()