servicediscovery.py

#!/usr/bin/env python3
"""
servicediscovery.py - OffSeq

CIDR-aware Nmap scanner that runs one subprocess per IP and writes XML output
per host into a results directory. Defaults to scanning all TCP ports with
moderate timing (T3) and runs 4 concurrent scans. UDP scans can be enabled
with rapid defaults.
"""

from __future__ import annotations

import argparse
import ipaddress
import os
import shutil
import subprocess
import sys
import time
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from pathlib import Path
from typing import Iterable, List


DEFAULT_THREADS = 4
DEFAULT_TIMING = 3
DEFAULT_UDP_TIMING = 4
DEFAULT_TOP_PORTS = 1000
DEFAULT_RESULTS_DIR = "results"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="servicediscovery.py",
        description=(
            "Scan CIDRs/IPs with Nmap using one subprocess per IP and save "
            "XML output into a results directory."
        ),
        epilog=(
            "Examples:\n"
            "  python3 servicediscovery.py 10.0.0.0/24\n"
            "  python3 servicediscovery.py 192.168.1.10 --test\n"
            "  python3 servicediscovery.py 10.0.0.0/24 --top 200 --threads 8\n"
            "  python3 servicediscovery.py 10.0.0.5 --timing 4\n"
            "  python3 servicediscovery.py 10.0.0.0/24 --udp\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "targets",
        nargs="+",
        help="Target CIDR(s) or IP address(es) (IPv4 or IPv6).",
    )
    parser.add_argument(
        "--threads",
        type=int,
        default=DEFAULT_THREADS,
        help=f"Number of parallel Nmap subprocesses (default {DEFAULT_THREADS}).",
    )
    parser.add_argument(
        "--timing",
        type=int,
        default=None,
        choices=range(0, 6),
        metavar="0-5",
        help=(
            "Nmap timing template (0-5). Default is 3 (normal), or 4 (rapid) "
            "when --udp is enabled and --timing is not provided."
        ),
    )
    parser.add_argument(
        "--results-dir",
        default=DEFAULT_RESULTS_DIR,
        help=f"Directory to store per-host XML results (default '{DEFAULT_RESULTS_DIR}').",
    )
    parser.add_argument(
        "--top",
        nargs="?",
        const=DEFAULT_TOP_PORTS,
        type=int,
        metavar="N",
        help=(
            "Scan top ports instead of all ports. Optionally specify the "
            f"number of top ports (default {DEFAULT_TOP_PORTS})."
        ),
    )
    parser.add_argument(
        "--test",
        action="store_true",
        help="Single-IP test scan using top ports only (requires exactly one IP).",
    )
    parser.add_argument(
        "--nmap-path",
        default=None,
        help="Optional path to the nmap binary (otherwise resolved from PATH).",
    )
    parser.add_argument(
        "--scan-type",
        choices=["syn", "connect"],
        default="syn",
        help="TCP scan type: syn (-sS, requires privileges) or connect (-sT).",
    )
    parser.add_argument(
        "--udp",
        action="store_true",
        help=(
            "Include UDP scan (-sU). When enabled without --timing, uses "
            "rapid timing (T4) by default."
        ),
    )
    parser.add_argument(
        "--ping",
        action="store_true",
        help="Enable host discovery (default is -Pn to treat all hosts as up).",
    )
    return parser.parse_args()


def resolve_nmap_path(user_path: str | None) -> str:
    if user_path:
        return user_path
    nmap_path = shutil.which("nmap")
    if not nmap_path:
        raise FileNotFoundError("nmap binary not found in PATH.")
    return nmap_path


def normalize_targets(targets: List[str]) -> List[ipaddress._BaseNetwork]:
    networks: List[ipaddress._BaseNetwork] = []
    for token in targets:
        try:
            if "/" in token:
                net = ipaddress.ip_network(token, strict=False)
            else:
                addr = ipaddress.ip_address(token)
                net = ipaddress.ip_network(f"{addr}/{addr.max_prefixlen}", strict=False)
            networks.append(net)
        except ValueError as exc:
            raise ValueError(f"Invalid target '{token}': {exc}") from exc
    return networks


def iter_addresses(networks: Iterable[ipaddress._BaseNetwork]) -> Iterable[ipaddress._BaseAddress]:
    for net in networks:
        for ip in net:
            yield ip


def sanitize_ip_for_filename(ip: ipaddress._BaseAddress) -> str:
    text = ip.compressed
    if ip.version == 6:
        text = text.replace(":", "_")
    return text


def build_nmap_command(
    ip: ipaddress._BaseAddress,
    nmap_path: str,
    timing: int,
    top_ports: int | None,
    scan_type: str,
    udp: bool,
    use_ping: bool,
    output_path: Path,
) -> List[str]:
    cmd: List[str] = [nmap_path]

    if scan_type == "syn":
        cmd.append("-sS")
    else:
        cmd.append("-sT")
    if udp:
        cmd.append("-sU")

    cmd.extend(["-sV", "-O"])
    if not use_ping:
        cmd.append("-Pn")

    cmd.append(f"-T{timing}")

    if top_ports is None:
        cmd.append("-p-")
    else:
        cmd.extend(["--top-ports", str(top_ports)])

    if ip.version == 6:
        cmd.append("-6")

    cmd.extend(["-oX", str(output_path), str(ip)])
    return cmd


def scan_one(
    ip: ipaddress._BaseAddress,
    nmap_path: str,
    timing: int,
    top_ports: int | None,
    scan_type: str,
    udp: bool,
    use_ping: bool,
    results_dir: Path,
) -> tuple[ipaddress._BaseAddress, int, float, str]:
    output_file = results_dir / f"{sanitize_ip_for_filename(ip)}.xml"
    cmd = build_nmap_command(
        ip=ip,
        nmap_path=nmap_path,
        timing=timing,
        top_ports=top_ports,
        scan_type=scan_type,
        udp=udp,
        use_ping=use_ping,
        output_path=output_file,
    )
    start = time.time()
    proc = subprocess.run(cmd, capture_output=True, text=True)
    duration = time.time() - start
    stderr = proc.stderr.strip()
    return ip, proc.returncode, duration, stderr


def main() -> int:
    args = parse_args()

    if args.threads < 1:
        print("Error: --threads must be >= 1.", file=sys.stderr)
        return 2

    if args.top is not None and args.top < 1:
        print("Error: --top must be >= 1.", file=sys.stderr)
        return 2

    try:
        nmap_path = resolve_nmap_path(args.nmap_path)
    except FileNotFoundError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 2

    if args.scan_type == "syn" and os.name != "nt":
        if hasattr(os, "geteuid") and os.geteuid() != 0:
            print(
                "Warning: SYN scans (-sS) usually require elevated privileges. "
                "Run as root or use --scan-type connect.",
                file=sys.stderr,
            )
    if args.udp and os.name != "nt":
        if hasattr(os, "geteuid") and os.geteuid() != 0:
            print(
                "Warning: UDP scans (-sU) can be limited without elevated "
                "privileges. Consider running as root for full capability.",
                file=sys.stderr,
            )

    try:
        networks = normalize_targets(args.targets)
    except ValueError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 2

    total_addresses = sum(net.num_addresses for net in networks)

    timing = args.timing if args.timing is not None else (
        DEFAULT_UDP_TIMING if args.udp else DEFAULT_TIMING
    )

    if args.test:
        if total_addresses != 1:
            print(
                "Error: --test requires exactly one IP address target.",
                file=sys.stderr,
            )
            return 2

    top_ports = None
    if args.test:
        top_ports = args.top if args.top is not None else DEFAULT_TOP_PORTS
    elif args.top is not None:
        top_ports = args.top

    results_dir = Path(args.results_dir)
    results_dir.mkdir(parents=True, exist_ok=True)

    print(
        f"Starting scan: targets={total_addresses}, "
        f"threads={args.threads}, timing=T{timing}, "
        f"ports={'top ' + str(top_ports) if top_ports else 'all'}, "
        f"udp={'on' if args.udp else 'off'}, "
        f"results_dir={results_dir}"
    )

    completed = 0
    failures = 0

    with ThreadPoolExecutor(max_workers=args.threads) as executor:
        futures = set()
        ip_iter = iter_addresses(networks)

        for ip in ip_iter:
            futures.add(
                executor.submit(
                    scan_one,
                    ip,
                    nmap_path,
                    timing,
                    top_ports,
                    args.scan_type,
                    args.udp,
                    args.ping,
                    results_dir,
                )
            )
            if len(futures) >= args.threads:
                done, futures = wait(futures, return_when=FIRST_COMPLETED)
                for fut in done:
                    ip, code, duration, stderr = fut.result()
                    completed += 1
                    if code != 0:
                        failures += 1
                        print(
                            f"[{completed}/{total_addresses}] {ip} FAILED "
                            f"(rc={code}, {duration:.1f}s): {stderr or 'no stderr'}"
                        )
                    else:
                        print(f"[{completed}/{total_addresses}] {ip} OK ({duration:.1f}s)")

        for fut in futures:
            ip, code, duration, stderr = fut.result()
            completed += 1
            if code != 0:
                failures += 1
                print(
                    f"[{completed}/{total_addresses}] {ip} FAILED "
                    f"(rc={code}, {duration:.1f}s): {stderr or 'no stderr'}"
                )
            else:
                print(f"[{completed}/{total_addresses}] {ip} OK ({duration:.1f}s)")

    print(f"Done. Completed={completed}, Failed={failures}, Results={results_dir}")
    return 0 if failures == 0 else 1


if __name__ == "__main__":
    sys.exit(main())