#!/usr/bin/env python3
"""
Author: Kieren Nicolas Lovell
Date: 26th November 2025
Purpose: Lightlevel scan of IP's to confirm open ports and identify banners.

elastic_scan.py  simple threaded port scanner for a list of IPs.

Examples:
    # Scan common ports on a few IPs
    python3 elastic_scan.py 54.216.18.65 3.250.26.12 54.216.48.67 --ports common

    # Scan all ports on all given IPs
    python3 elastic_scan.py 54.216.18.65 3.250.26.12 --ports all

    # Scan a custom list/range of ports
    python3 elastic_scan.py 54.216.18.65 --ports 22,80,443,8080-8090

    # Use CIDR ranges and commas/spaces
    python3 elastic_scan.py 207.228.129.226/30,87.228.231.46/30 195.244.216.10/29 91.74.32.202 --ports all
"""

import argparse
import socket
import ipaddress
from concurrent.futures import ThreadPoolExecutor, as_completed

COMMON_PORTS = list(range(1, 1025)) + [
    1433, 1521, 2049, 3306, 3389, 5432, 5900, 6379, 8080, 8443, 9200, 11211
]


def parse_port_spec(spec: str):
    """Parse a port spec like '22,80,443,8000-8100' into a sorted list."""
    ports = set()
    for part in spec.split(","):
        part = part.strip()
        if not part:
            continue
        if "-" in part:
            start_s, end_s = part.split("-", 1)
            start, end = int(start_s), int(end_s)
            if start > end:
                start, end = end, start
            ports.update(range(start, end + 1))
        else:
            ports.add(int(part))
    return sorted(p for p in ports if 1 <= p <= 65535)


def make_port_list(mode: str):
    mode = mode.lower()
    if mode == "common":
        return COMMON_PORTS
    if mode == "all":
        return list(range(1, 65536))
    # Otherwise treat as explicit spec
    return parse_port_spec(mode)


def expand_ip_specs(ip_specs: list[str]) -> list[str]:
    """
    Expand CLI IP specs (which may contain commas and CIDR) into a flat list
    of unique IPv4 addresses as strings.
    """
    result: list[str] = []
    seen: set[str] = set()

    for spec in ip_specs:
        parts = [p.strip() for p in spec.split(",") if p.strip()]
        for part in parts:
            try:
                if "/" in part:
                    net = ipaddress.ip_network(part, strict=False)
                    for host in net.hosts():
                        ip_str = str(host)
                        if ip_str not in seen:
                            result.append(ip_str)
                            seen.add(ip_str)
                else:
                    ip_obj = ipaddress.ip_address(part)
                    ip_str = str(ip_obj)
                    if ip_str not in seen:
                        result.append(ip_str)
                        seen.add(ip_str)
            except ValueError as e:
                raise SystemExit(f"[!] Invalid IP or CIDR specification: '{part}' ({e})")

    return result


def scan_single_port(ip: str, port: int, timeout: float) -> int | None:
    """Return port if open, otherwise None."""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(timeout)
            result = s.connect_ex((ip, port))
            if result == 0:
                return port
    except OSError:
        pass
    return None


def scan_ip(ip: str, ports: list[int], timeout: float, max_workers: int):
    """Scan a set of ports on a single IP and return a sorted list of open ports."""
    open_ports: list[int] = []
    workers = min(max_workers, len(ports)) or 1
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(scan_single_port, ip, port, timeout): port
            for port in ports
        }
        for future in as_completed(futures):
            port = future.result()
            if port is not None:
                open_ports.append(port)
    return sorted(open_ports)


def grab_banner(ip: str, port: int, timeout: float) -> str:
    """
    Simple banner grab: connect, send a small protocol-specific probe if we know one,
    then read up to 1024 bytes. Returns a decoded string (may be empty).
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(timeout)
            s.connect((ip, port))

            # HTTP(S)-style ports  send a minimal HEAD
            if port in (80, 81, 443, 8000, 8080, 8443, 8888):
                payload = (
                    b"HEAD / HTTP/1.0\r\n"
                    b"Host: " + ip.encode(errors="ignore") + b"\r\n"
                    b"\r\n"
                )
                s.sendall(payload)

            # SIP (TCP 5060/5061)  send a basic OPTIONS request
            elif port in (5060, 5061):
                payload = (
                    "OPTIONS sip:example.com SIP/2.0\r\n"
                    "Via: SIP/2.0/TCP scanner;branch=z9hG4bK-1\r\n"
                    "From: \"scanner\" <sip:scanner@example.com>;tag=1\r\n"
                    "To: <sip:scanner@example.com>\r\n"
                    "Call-ID: 1@scanner\r\n"
                    "CSeq: 1 OPTIONS\r\n"
                    "Content-Length: 0\r\n"
                    "\r\n"
                ).encode()
                s.sendall(payload)

            else:
                # Generic nudge  sometimes triggers a banner
                try:
                    s.sendall(b"\r\n")
                except OSError:
                    pass

            data = s.recv(1024)
            if not data:
                return ""
            return data.decode(errors="replace").strip()
    except OSError:
        return ""


def guess_service(port: int, banner: str) -> str | None:
    """
    Lightweight heuristic to guess what is running on a port
    based on port number and banner contents.
    """
    b = (banner or "").lower()

    # Port-based hints first
    if port == 22 or "ssh" in b:
        return "ssh"
    if port in (80, 81, 443, 8000, 8080, 8443, 8888) or "http/" in b or "server:" in b:
        return "http"
    if port in (25, 587, 465) or "smtp" in b:
        return "smtp"
    if port in (110, 995) or "pop3" in b:
        return "pop3"
    if port in (143, 993) or "imap" in b:
        return "imap"
    if port == 53 or "dns" in b:
        return "dns"
    if port == 21 or "ftp" in b:
        return "ftp"
    if port == 3306 or "mysql" in b:
        return "mysql"
    if port == 5432 or "postgres" in b or "postgresql" in b:
        return "postgresql"
    if port == 6379 or "redis" in b:
        return "redis"
    if port == 3389 or "rdp" in b or "remote desktop" in b:
        return "rdp"
    if port in (5060, 5061) or "sip/" in b or "via: sip" in b:
        return "sip"

    # Generic protocol markers
    if "ssl" in b or "tls" in b:
        return "ssl/tls"
    if "soap" in b or "xml" in b:
        return "web-service"

    return None


def collect_banners(ip: str, open_ports: list[int], timeout: float) -> dict[int, str]:
    """
    For each open port, attempt to grab a banner.
    Returns {port: banner} for ports with non-empty banners.
    """
    banners: dict[int, str] = {}
    # Give banner grabbing a bit more time than the scan itself
    banner_timeout = timeout * 2 if timeout > 0 else 1.0

    for port in open_ports:
        banner = grab_banner(ip, port, banner_timeout)
        if banner:
            if len(banner) > 200:
                banner = banner[:197] + "..."
            banners[port] = banner
    return banners


def main():
    parser = argparse.ArgumentParser(
        description="Simple TCP port scanner for IPs and CIDR ranges, with basic banner grabbing."
    )
    parser.add_argument(
        "ips",
        nargs="+",
        help=(
            "IP addresses and/or CIDR ranges to scan. "
            "Supports spaces and/or commas, e.g.: "
            "'207.228.129.226/30,87.228.231.46/30 195.244.216.10/29 91.74.32.202'"
        ),
    )
    parser.add_argument(
        "--ports",
        default="common",
        help=(
            "Port selection: 'common', 'all', or custom spec "
            "like '22,80,443,8000-8100'. Default: common"
        ),
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=0.5,
        help="Socket timeout per connection in seconds (default: 0.5)",
    )
    parser.add_argument(
        "--threads",
        type=int,
        default=200,
        help="Max threads per host (default: 200)",
    )
    parser.add_argument(
        "--no-banner",
        action="store_true",
        help="Disable banner grabbing / service guessing (scan only).",
    )

    args = parser.parse_args()
    ports = make_port_list(args.ports)
    all_ips = expand_ip_specs(args.ips)

    if not all_ips:
        raise SystemExit("[!] No valid IPs to scan after expanding input.")

    print(f"[+] Expanded to {len(all_ips)} host(s), {len(ports)} port(s) each.")
    print(f"[+] Mode: {args.ports}, timeout={args.timeout}s, threads={args.threads}")
    if not args.no_banner:
        print("[+] Banner grabbing: enabled\n")
    else:
        print("[+] Banner grabbing: disabled\n")

    for ip in all_ips:
        print(f"=== {ip} ===")
        open_ports = scan_ip(ip, ports, args.timeout, args.threads)
        if not open_ports:
            print("No open ports found (in this port set).\n")
            continue

        if args.no_banner:
            print("Open ports:")
            print("  " + ", ".join(str(p) for p in open_ports))
            print()
            continue

        banners = collect_banners(ip, open_ports, args.timeout)

        print("Open ports:")
        for port in open_ports:
            banner = banners.get(port, "")
            service = guess_service(port, banner)

            if banner:
                if service:
                    print(f"  {port} ({service}): {banner}")
                else:
                    print(f"  {port}: {banner}")
            else:
                # No banner body, but we still may have a decent guess from the port
                if service:
                    print(f"  {port} ({service}): (open, no banner)")
                else:
                    print(f"  {port}: (open, no banner)")
        print()


if __name__ == "__main__":
    main()