package unit;

import java.util.*;

/**
 * nmap-0001 — CWE-407: O(n²) port membership test in version-detection hot path
 *
 * Models nmap service_scan.cc ServiceProbe::portIsProbable():
 *   Slow: find(portv->begin(), portv->end(), portno)  — O(K) per probe per query
 *   Fast: unordered_set<u16>::count(portno)           — O(1) per probe per query
 *
 * Hot path: nextProbe() iterates P=187 probes, each calling portIsProbable().
 * With K ports per probe, each service scan step costs O(P*K).
 * For large port lists (up to 32,771 expanded) this dominates -sV runtime.
 */
public class NmapPortMembershipTest {

    // --- SLOW: vector linear scan (defective) ---
    static boolean portIsProbableSlow(List<Integer> portv, int portno) {
        return portv.contains(portno); // O(K) linear
    }

    static int runSlowScan(List<List<Integer>> probes, int[] portsToCheck) {
        int ops = 0;
        for (int port : portsToCheck) {
            for (List<Integer> probe : probes) {
                // simulate: every portIsProbable check scans the whole vector
                for (int p : probe) {
                    ops++;
                    if (p == port) break;
                }
            }
        }
        return ops;
    }

    // --- FAST: unordered_set O(1) lookup (fixed) ---
    static boolean portIsProbableFast(Set<Integer> portSet, int portno) {
        return portSet.contains(portno); // O(1) hash lookup
    }

    static int runFastScan(List<Set<Integer>> probeSets, int[] portsToCheck) {
        int ops = 0;
        for (int port : portsToCheck) {
            for (Set<Integer> probeSet : probeSets) {
                ops++; // single hash lookup per probe
                probeSet.contains(port);
            }
        }
        return ops;
    }

    // Build probe port lists: N probes, each with K ports starting at offset
    static List<List<Integer>> buildSlowProbes(int N, int K) {
        List<List<Integer>> probes = new ArrayList<>();
        for (int i = 0; i < N; i++) {
            List<Integer> ports = new ArrayList<>();
            for (int j = 0; j < K; j++) {
                ports.add(1024 + (i * K + j) % 60000);
            }
            probes.add(ports);
        }
        return probes;
    }

    static List<Set<Integer>> buildFastProbes(int N, int K) {
        List<Set<Integer>> probes = new ArrayList<>();
        for (int i = 0; i < N; i++) {
            Set<Integer> ports = new HashSet<>();
            for (int j = 0; j < K; j++) {
                ports.add(1024 + (i * K + j) % 60000);
            }
            probes.add(ports);
        }
        return probes;
    }

    public static void main(String[] args) {
        // N = number of probes (nmap has 187; use 200 for round number)
        // K = ports per probe (worst case: 32,771; use 1000 as representative large probe)
        // Q = number of port queries (simulates one -sV scan pass)
        final int N = 200;
        final int K = 1000;
        final int Q = 50;

        int[] portsToCheck = new int[Q];
        Random rng = new Random(42);
        for (int i = 0; i < Q; i++) {
            portsToCheck[i] = rng.nextInt(65535);
        }

        List<List<Integer>> slowProbes = buildSlowProbes(N, K);
        List<Set<Integer>> fastProbes = buildFastProbes(N, K);

        // Correctness check: same port should match (or not) in both versions
        // Use probe 0 with a known port
        int knownPort = slowProbes.get(0).get(0);
        int absentPort = 9; // not in any probe by construction (base 1024)
        assert portIsProbableSlow(slowProbes.get(0), knownPort)  : "slow: should find known port";
        assert !portIsProbableSlow(slowProbes.get(0), absentPort): "slow: should miss absent port";
        assert portIsProbableFast(fastProbes.get(0), knownPort)  : "fast: should find known port";
        assert !portIsProbableFast(fastProbes.get(0), absentPort): "fast: should miss absent port";

        // Operation count comparison
        int slowOps = runSlowScan(slowProbes, portsToCheck);
        int fastOps = runFastScan(fastProbes, portsToCheck);

        // Slow: up to Q * N * K operations worst case = 50 * 200 * 1000 = 10,000,000
        // Fast: exactly Q * N operations = 50 * 200 = 10,000
        long expectedSlowMin = (long) Q * N;        // at least one op per probe (found at start)
        long expectedSlowMax = (long) Q * N * K;    // worst case: never found
        long expectedFast    = (long) Q * N;        // always exactly one hash lookup

        assert slowOps > fastOps : "slow should cost more ops than fast, got slow=" + slowOps + " fast=" + fastOps;
        assert fastOps == (long) Q * N : "fast ops should be Q*N=" + (Q*N) + " got " + fastOps;

        double speedup = (double) slowOps / fastOps;

        System.out.println("nmap-0001 CWE-407: portIsProbable O(n) vector find vs O(1) hash set");
        System.out.println("  N (probes) = " + N + ", K (ports/probe) = " + K + ", Q (queries) = " + Q);
        System.out.println("  Slow ops (vector linear scan): " + slowOps);
        System.out.println("  Fast ops (hash set lookup):    " + fastOps);
        System.out.printf ("  Speedup: %.0fx%n", speedup);
        System.out.println();

        // serviceIsPossible: separate linear scan over detectedServices
        // Slow: vector<const char*> with strcmp loop, O(D) per call
        // Fast: unordered_set<string>, O(1) per call
        // D = detected services per probe (up to ~10 soft matches)
        int D = 10;
        int probeCount = 200;

        List<String> slowServices = new ArrayList<>();
        Set<String> fastServices = new HashSet<>();
        String[] serviceNames = {"http","ssh","ftp","smtp","pop3","imap","mysql","postgresql","redis","memcached"};
        for (String s : serviceNames) { slowServices.add(s); fastServices.add(s); }

        int slowSvcOps = 0, fastSvcOps = 0;
        String target = "redis"; // appears near end of list
        for (int i = 0; i < probeCount; i++) {
            // slow: scan the list with strcmp
            for (String s : slowServices) {
                slowSvcOps++;
                if (s.equals(target)) break;
            }
            // fast: single hash lookup
            fastSvcOps++;
            fastServices.contains(target);
        }

        assert slowSvcOps > fastSvcOps : "serviceIsPossible: slow should cost more";

        double svcSpeedup = (double) slowSvcOps / fastSvcOps;
        System.out.println("nmap-0001 CWE-407: serviceIsPossible O(n) strcmp loop vs O(1) hash set");
        System.out.println("  D (services/probe) = " + D + ", probes = " + probeCount);
        System.out.println("  Slow ops (linear strcmp): " + slowSvcOps);
        System.out.println("  Fast ops (hash set):      " + fastSvcOps);
        System.out.printf ("  Speedup: %.0fx%n", svcSpeedup);
        System.out.println();

        System.out.println("2/2 PASS");
    }
}