chain_engine.py

# ============================================================
#  KnockSense — Chain Engine
#  Handles multi-stage port knocking with timing precision
# ============================================================

import socket, threading, time
from modules.colors import *

def knock_tcp(host, port):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(1)
        s.connect_ex((host, port))
        s.close()
    except: pass

def knock_udp(host, port):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(0.5)
        s.sendto(b'\x00', (host, port))
        s.close()
    except: pass

def knock_simultaneous(host, mixed_ports):
    """
    Knock all ports at exactly the same time.
    mixed_ports: {port: 'tcp'/'udp'}
    """
    def worker(port, proto):
        if proto == 'udp':
            knock_udp(host, port)
        else:
            knock_tcp(host, port)
    threads = [
        threading.Thread(target=worker, args=(p, proto))
        for p, proto in mixed_ports.items()
    ]
    for t in threads: t.start()
    for t in threads: t.join()

def check_port(host, port, timeout=0.5):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        r = s.connect_ex((host, port))
        s.close()
        return r == 0
    except: return False

def quick_scan(host, ports, timeout=0.3):
    """Scan specific ports fast."""
    opened = set()
    lock = threading.Lock()
    def check(p):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(timeout)
            r = s.connect_ex((host, p))
            s.close()
            if r == 0:
                with lock: opened.add(p)
        except: pass
    threads = [threading.Thread(target=check, args=(p,)) for p in ports]
    for t in threads: t.start()
    for t in threads: t.join()
    return opened

def _crawl_and_store(host, port, all_opened, pcap_list, data_list):
    """Immediately crawl an opened port and collect PCAPs."""
    import os, sys
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    try:
        from modules.extractor import _fetch_url, find_pcap_links, _download_file
        dl_dir = "/tmp/knocksense_downloads"
        os.makedirs(dl_dir, exist_ok=True)
        proto = "https" if port in [443, 8443] else "http"
        # Try root path first
        for path in ["/", "/iamcornholio/", "/spanishfly/", "/burgerworld/"]:
            url = proto + "://" + host + ":" + str(port) + path
            page = _fetch_url(url, timeout=3)
            if not page: continue
            for pcap_url in find_pcap_links(page, url):
                pcap_name = pcap_url.split("/")[-1]
                save_path = os.path.join(dl_dir, pcap_name)
                counter = 1
                base = pcap_name.rsplit(".", 1)
                while os.path.exists(save_path):
                    save_path = os.path.join(dl_dir, base[0] + "_" + str(counter) + "." + base[1])
                    counter += 1
                size = _download_file(pcap_url, save_path)
                if size > 100:
                    print(f"  \033[95m[>>>]\033[0m Downloaded: {pcap_name} ({size} bytes)")
                    pcap_list.append(save_path)
    except Exception as e:
        pass

def execute_chain(host, stages, baseline, verbose=True):
    """
    Execute multi-stage knock chain with timing precision.

    stages: list of dicts:
    [
        {
            'sequence': [7000, 8000, 9000],
            'mixed_ports': {7000:'tcp', 8000:'tcp', 9000:'tcp'},
            'opens_port': 8888,  # expected port to open (optional)
        },
        {
            'sequence': [21, 22, 80, 8080],
            'mixed_ports': {21:'tcp', 22:'udp', 80:'tcp', 8080:'tcp'},
        }
    ]
    """
    CHECK_PORTS = list(range(1, 1025)) + [
        1337, 2222, 3333, 4444, 5555, 6666, 7777,
        8080, 8443, 8888, 9000, 9090, 9999, 10000, 31337
    ]

    all_opened = set()
    stage_pcaps = []
    stage_data = []

    for i, stage in enumerate(stages):
        sequence = stage.get('sequence') or stage.get('ports', [])
        mixed = stage.get('mixed_ports', {p: 'tcp' for p in sequence})
        expected = stage.get('opens_port')

        if verbose:
            info(f"Stage {i+1}: {Y}{' → '.join(map(str, sequence))}{RESET}")

        # For stages after first: re-knock ALL previous stages first
        # to ensure gated ports stay open
        if i > 0:
            if verbose:
                info(f"Re-knocking previous {i} stage(s) to keep gates open...")
            for prev_stage in stages[:i]:
                prev_seq = prev_stage.get('sequence') or prev_stage.get('ports', [])
                prev_mixed = prev_stage.get('mixed_ports',
                             {p: 'tcp' for p in prev_seq})
                knock_simultaneous(host, prev_mixed)
                time.sleep(0.3)

        # Now knock this stage
        knock_simultaneous(host, mixed)

        if verbose:
            parts = [str(p)+':'+proto.upper() for p,proto in mixed.items()]
            print(f"  {G}[chain knock]{RESET} {' → '.join(parts)} ✓")

        # Check result
        time.sleep(1.0)
        after = quick_scan(host, CHECK_PORTS, timeout=0.3)
        newly_opened = after - baseline - all_opened

        if newly_opened:
            for p in sorted(newly_opened):
                found(f"Stage {i+1} opened port: {p} ✅")
            all_opened |= newly_opened
        elif expected and not check_port(host, expected):
            warn(f"Stage {i+1}: expected port {expected} not open yet")
            # Try once more
            knock_simultaneous(host, mixed)
            time.sleep(0.8)
            after = quick_scan(host, CHECK_PORTS, timeout=0.3)
            newly_opened = after - baseline - all_opened
            if newly_opened:
                for p in sorted(newly_opened):
                    found(f"Stage {i+1} opened port: {p} ✅ (retry)")
                all_opened |= newly_opened

    return all_opened, stage_pcaps


def auto_chain(host, pcap_sequences, baseline, verbose=True):
    """
    Automatically build and execute a chain from multiple PCAP sequences.
    pcap_sequences: list of sequence dicts from analyze_pcap_protocols()
    """
    if not pcap_sequences:
        return set()

    # Build stages from sequences
    stages = []
    for seq_data in pcap_sequences:
        ports = seq_data['ports']
        mixed = seq_data.get('mixed_ports', {p: 'tcp' for p in ports})
        # If no mixed_ports info, apply heuristics
        if not mixed or all(v == 'tcp' for v in mixed.values()):
            # Port 22 is commonly UDP in CTF knocks
            UDP_HINTS = {22, 53, 69}
            mixed = {p: 'udp' if p in UDP_HINTS else 'tcp' for p in ports}
        stages.append({
            'sequence': ports,
            'mixed_ports': mixed,
        })

    return execute_chain(host, stages, baseline, verbose)