package unit;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * strongswan-0001: CWE-407 O(P²×T×A²) nested linear scans in proposal_select / select_algo
 *
 * Models the strongSwan proposal negotiation:
 *   SLOW: nested ArrayList scans for per-transform algorithm matching (O(A₁×A₂) per type)
 *         inside an outer O(P_c × P_s) proposal loop.
 *   FAST: HashMap-indexed algorithm set per transform type → O(A₁+A₂)
 *
 * No JUnit — compile and run standalone.
 */
public class StrongSwanProposalSelectTest {

    // ------------------------------------------------------------------
    // Data model: a "proposal" has transform types, each with a list of algs
    // Each alg is an int (algorithm id | key_size packed as long)
    // ------------------------------------------------------------------
    static class Proposal {
        // transform_type -> list of (alg_id << 16 | key_size) longs
        List<List<Long>> transforms; // indexed by transform type 0..T-1

        Proposal(int numTypes) {
            transforms = new ArrayList<>();
            for (int i = 0; i < numTypes; i++) transforms.add(new ArrayList<>());
        }

        void addAlg(int type, long algKey) {
            transforms.get(type).add(algKey);
        }
    }

    // ------------------------------------------------------------------
    // SLOW: select_algo O(A₁ × A₂) — mirrors the C nested enumerator loop
    //   Returns op count.
    // ------------------------------------------------------------------
    static long slowSelectAlgo(List<Long> localAlgs, List<Long> remoteAlgs) {
        long ops = 0;
        for (Long alg1 : localAlgs) {
            for (Long alg2 : remoteAlgs) {
                ops++;
                if (alg1.equals(alg2)) {
                    return ops; // found — mirrors early break
                }
            }
        }
        return ops; // no match
    }

    // SLOW: select_algos — iterate all T transform types
    static long slowSelectAlgos(Proposal local, Proposal remote) {
        long ops = 0;
        int T = local.transforms.size();
        for (int type = 0; type < T; type++) {
            ops += slowSelectAlgo(local.transforms.get(type), remote.transforms.get(type));
        }
        return ops;
    }

    // SLOW: proposal_select — O(P_c × P_s) outer loop
    static long slowProposalSelect(List<Proposal> configured, List<Proposal> supplied) {
        long ops = 0;
        for (Proposal local : configured) {
            for (Proposal remote : supplied) {
                ops += slowSelectAlgos(local, remote);
            }
        }
        return ops;
    }

    // ------------------------------------------------------------------
    // FAST: build HashMap per transform type from remote proposal, then O(1) lookup
    // ------------------------------------------------------------------
    static long fastSelectAlgo(List<Long> localAlgs, List<Long> remoteAlgs) {
        HashMap<Long, Boolean> remoteSet = new HashMap<>();
        long ops = 0;
        for (Long alg2 : remoteAlgs) {
            ops++;
            remoteSet.put(alg2, true);
        }
        for (Long alg1 : localAlgs) {
            ops++;
            if (remoteSet.containsKey(alg1)) {
                return ops; // found
            }
        }
        return ops;
    }

    static long fastSelectAlgos(Proposal local, Proposal remote) {
        long ops = 0;
        int T = local.transforms.size();
        for (int type = 0; type < T; type++) {
            ops += fastSelectAlgo(local.transforms.get(type), remote.transforms.get(type));
        }
        return ops;
    }

    static long fastProposalSelect(List<Proposal> configured, List<Proposal> supplied) {
        long ops = 0;
        for (Proposal local : configured) {
            for (Proposal remote : supplied) {
                ops += fastSelectAlgos(local, remote);
            }
        }
        return ops;
    }

    // ------------------------------------------------------------------
    // Helpers: build proposal lists
    // ------------------------------------------------------------------

    /**
     * Build a list of P proposals, each with T transform types,
     * each type having A algorithms. Algorithms are unique per proposal
     * and offset so that no match occurs until the last comparison
     * (worst-case for the slow path).
     */
    static List<Proposal> makeProposals(int P, int T, int A, int baseOffset) {
        List<Proposal> proposals = new ArrayList<>();
        for (int p = 0; p < P; p++) {
            Proposal prop = new Proposal(T);
            for (int t = 0; t < T; t++) {
                for (int a = 0; a < A; a++) {
                    // unique algorithm id per (proposal, type, alg)
                    long algKey = (long)(baseOffset + p * T * A + t * A + a) << 16;
                    prop.addAlg(t, algKey);
                }
            }
            proposals.add(prop);
        }
        return proposals;
    }

    /**
     * Build supplied proposals that match on the LAST algorithm of LAST
     * proposal to force worst-case O(P²×T×A²).
     */
    static List<Proposal> makeMatchingSupplied(List<Proposal> configured, int T, int A) {
        List<Proposal> supplied = new ArrayList<>();
        // clone configured proposals, put the matching alg last in each type
        for (Proposal cfg : configured) {
            Proposal sup = new Proposal(T);
            for (int t = 0; t < T; t++) {
                List<Long> cfgAlgs = cfg.transforms.get(t);
                // add non-matching alg first, then the matching alg last
                sup.addAlg(t, Long.MAX_VALUE - t); // no-match filler
                sup.addAlg(t, cfgAlgs.get(cfgAlgs.size() - 1)); // last cfg alg
            }
            supplied.add(sup);
        }
        return supplied;
    }

    public static void main(String[] args) {
        int tests = 0, passed = 0;

        // --- correctness: tiny case ---
        {
            tests++;
            int P = 2, T = 2, A = 3;
            List<Proposal> configured = makeProposals(P, T, A, 0);
            List<Proposal> supplied = makeMatchingSupplied(configured, T, A);

            long slowOps = slowProposalSelect(configured, supplied);
            long fastOps = fastProposalSelect(configured, supplied);

            // both should find a match (ops > 0); fast should use fewer ops
            if (slowOps > 0 && fastOps > 0 && fastOps <= slowOps) {
                passed++;
                System.out.println("PASS correctness P=" + P + " T=" + T + " A=" + A
                    + " slow_ops=" + slowOps + " fast_ops=" + fastOps);
            } else {
                System.out.println("FAIL correctness P=" + P + " T=" + T + " A=" + A
                    + " slow_ops=" + slowOps + " fast_ops=" + fastOps);
            }
        }

        // --- no-match worst case ---
        {
            tests++;
            int P = 5, T = 3, A = 5;
            List<Proposal> configured = makeProposals(P, T, A, 0);
            // supplied has completely different algs — no match at all
            List<Proposal> supplied = makeProposals(P, T, A, 10000);

            long slowOps = slowProposalSelect(configured, supplied);
            long fastOps = fastProposalSelect(configured, supplied);

            double ratio = (double) slowOps / fastOps;
            boolean ok = ratio >= 2.0;
            if (ok) {
                passed++;
                System.out.printf("PASS no-match P=%d T=%d A=%d slow_ops=%d fast_ops=%d ratio=%.1fx%n",
                    P, T, A, slowOps, fastOps, ratio);
            } else {
                System.out.printf("FAIL no-match P=%d T=%d A=%d slow_ops=%d fast_ops=%d ratio=%.1fx (need >=2x)%n",
                    P, T, A, slowOps, fastOps, ratio);
            }
        }

        // --- speedup benchmarks ---
        int[][] benchParams = {
            // {P, T, A}  — worst case: no match until last comparison
            {10, 5, 8},
            {20, 5, 8},
            {30, 7, 10},
            {50, 7, 10},
        };
        for (int[] param : benchParams) {
            tests++;
            int P = param[0], T = param[1], A = param[2];
            List<Proposal> configured = makeProposals(P, T, A, 0);
            // no-match supplied: completely different algorithm ids
            List<Proposal> supplied = makeProposals(P, T, A, 100000);

            long slowOps = slowProposalSelect(configured, supplied);
            long fastOps = fastProposalSelect(configured, supplied);

            double ratio = (double) slowOps / fastOps;
            // At A=8 we expect ~A ratio for select_algo alone; with P² multiplied, ratio should be >=3x
            boolean ok = ratio >= 3.0;
            if (ok) {
                passed++;
                System.out.printf("PASS speedup P=%d T=%d A=%d slow_ops=%d fast_ops=%d ratio=%.1fx%n",
                    P, T, A, slowOps, fastOps, ratio);
            } else {
                System.out.printf("FAIL speedup P=%d T=%d A=%d slow_ops=%d fast_ops=%d ratio=%.1fx (need >=3x)%n",
                    P, T, A, slowOps, fastOps, ratio);
            }
        }

        // --- large-scale attack simulation ---
        // Attacker sends 50 non-matching proposals, server has 30 proposals
        // This models a DoS amplification scenario
        {
            tests++;
            int Pc = 30, Ps = 50, T = 5, A = 8;
            List<Proposal> configured = makeProposals(Pc, T, A, 0);
            List<Proposal> supplied = makeProposals(Ps, T, A, 100000);

            long slowOps = slowProposalSelect(configured, supplied);
            long fastOps = fastProposalSelect(configured, supplied);

            double ratio = (double) slowOps / fastOps;
            // With A=8: slow does A² ops per transform (64), fast does 2A (16) → ratio ~= A/2 = 4x
            // The key point is the A² exponent is eliminated; require ratio >= 3x
            boolean ok = ratio >= 3.0;
            if (ok) {
                passed++;
                System.out.printf("PASS dos-scenario Pc=%d Ps=%d T=%d A=%d slow_ops=%d fast_ops=%d ratio=%.1fx%n",
                    Pc, Ps, T, A, slowOps, fastOps, ratio);
            } else {
                System.out.printf("FAIL dos-scenario Pc=%d Ps=%d T=%d A=%d slow_ops=%d fast_ops=%d ratio=%.1fx (need >=3x)%n",
                    Pc, Ps, T, A, slowOps, fastOps, ratio);
            }
        }

        System.out.println("\n" + passed + "/" + tests + " PASS");
        if (passed != tests) System.exit(1);
    }
}