from bbot.errors import HttpCompareError
from .base import BaseLightfuzz

import urllib.parse


class cmdi(BaseLightfuzz):
    """
    Detects command injection vulnerabilities.

    Techniques:

    * Echo Canary Detection:
       - Injects command delimiters (;, &&, ||, &, |) along with an echo command
       - Checks if the echoed canary appears in the response without the "echo" itself
       - Uses a false positive probe to validate findings

    * Arithmetic Canary Detection (confirmation cascade):
       - When an echo canary match is found, two follow-up probes try to prove
         the response was generated by real command execution rather than
         textual reflection of probe bytes:
           - `echo $((A*B))` — POSIX arithmetic expansion (bash/sh/dash/zsh/ksh)
           - `set /A A*B`   — Windows cmd.exe native arithmetic
         The product A*B never appears in either probe's literal bytes, so
         parser-error or template reflections cannot produce it. A match
         upgrades the finding to HIGH confidence and labels the shell family.
       - If neither arithmetic probe confirms execution, the generic match
         is emitted as a "Possible Parameter Reflection" finding (LOW severity)
         instead of a cmdi finding — this is honest about what we observed
         (the value is reflected into the response) while still surfacing the
         endpoint for triage, since reflected values often point at adjacent
         vulns the other submodules missed (e.g. SQLi).

    * Blind Command Injection:
       - Injects nslookup commands with unique subdomain tags
       - Detects command execution through DNS resolution via Interactsh
    """

    friendly_name = "Command Injection"
    uses_interactsh = True

    async def fuzz(self):
        cookies = self.event.data.get(
            "assigned_cookies", {}
        )  # Retrieve assigned cookies from WEB_PARAMETER event data, if present
        probe_value = self.incoming_probe_value()

        canary = self.lightfuzz.helpers.rand_string(10, numeric_only=True)
        # Arithmetic canary: the multiplicands are in the probe, but their
        # product is NOT — so a response containing the product is evidence
        # of arithmetic expansion performed by a real shell. Reject factors
        # that evaluate to 0; otherwise the canary degenerates to "0" and
        # matches virtually every HTML response (critical-severity FP).
        while True:
            arith_a = self.lightfuzz.helpers.rand_string(5, numeric_only=True)
            arith_b = self.lightfuzz.helpers.rand_string(5, numeric_only=True)
            if int(arith_a) and int(arith_b) and arith_a[0] != "0" and arith_b[0] != "0":
                break
        arith_canary = str(int(arith_a) * int(arith_b))
        http_compare = self.compare_baseline(
            self.event.data["type"], probe_value, cookies
        )  # Initialize the http_compare object and establish a baseline HTTP response

        cmdi_probe_strings = [
            "AAAA",  # False positive probe
            ";",
            "&&",
            "||",
            "&",
            "|",
        ]

        linux_detections = []  # POSIX arithmetic canary matched → HIGH
        windows_detections = []  # cmd.exe arithmetic canary matched → HIGH
        reflection_detections = []  # generic matched but neither shell proved execution
        for p in cmdi_probe_strings:
            try:
                generic_probe_str = f"{probe_value}{p} echo {canary} {p}"
                if self.event.data["type"] == "GETPARAM":
                    generic_probe_str = urllib.parse.quote(generic_probe_str.encode(), safe="")

                generic_probe = await self.compare_probe(
                    http_compare, self.event.data["type"], generic_probe_str, cookies, skip_urlencoding=True
                )

                generic_match = (
                    generic_probe[3] is not None
                    and canary in generic_probe[3].text
                    and "echo" not in generic_probe[3].text
                )
                if not generic_match:
                    continue
                if p == "AAAA":
                    self.warning(
                        f"False Postive Probe appears to have been triggered for {self.event.url}, aborting remaining detection"
                    )
                    return

                # Cascade: try POSIX arithmetic first, then Windows cmd
                # arithmetic. The product value is never in either probe
                # literal, so only real shell execution can place it in the
                # response. Either confirmation upgrades to HIGH.
                raw_marker = f"{arith_a}*{arith_b}"
                linux_match = await self._arith_confirm(
                    http_compare,
                    cookies,
                    f"{probe_value}{p} echo $(({arith_a}*{arith_b})) {p}",
                    arith_canary,
                    raw_marker,
                    p,
                    "linux",
                )
                if linux_match:
                    linux_detections.append(p)
                    continue
                windows_match = await self._arith_confirm(
                    http_compare,
                    cookies,
                    f"{probe_value}{p} set /A {arith_a}*{arith_b} {p}",
                    arith_canary,
                    raw_marker,
                    p,
                    "windows",
                )
                if windows_match:
                    windows_detections.append(p)
                    continue
                # Generic fired, neither shell confirmed — reflection.
                reflection_detections.append(p)
                self.debug(
                    f"echo canary reflected without arithmetic confirmation for [{p}] — "
                    f"treating as parameter reflection, not cmdi"
                )
            except HttpCompareError as e:
                self.debug(e)
                continue
        if linux_detections:
            self.results.append(
                {
                    "name": "Possible Command Injection",
                    "severity": "CRITICAL",
                    "confidence": "HIGH",
                    "description": f"POSSIBLE OS Command Injection. {self.metadata()} Detection Method: [arithmetic canary (POSIX)] CMD Probe Delimeters: [{' '.join(linux_detections)}]",
                }
            )
        if windows_detections:
            self.results.append(
                {
                    "name": "Possible Command Injection",
                    "severity": "CRITICAL",
                    "confidence": "HIGH",
                    "description": f"POSSIBLE OS Command Injection. {self.metadata()} Detection Method: [arithmetic canary (cmd)] CMD Probe Delimeters: [{' '.join(windows_detections)}]",
                }
            )
        if reflection_detections:
            self.results.append(
                {
                    "name": "Possible Parameter Reflection",
                    "severity": "LOW",
                    "confidence": "LOW",
                    "description": f"Parameter value reflected without confirmed shell execution. May indicate adjacent reflection-based vulns. {self.metadata()} Detection Method: [echo canary, no shell confirmation] CMD Probe Delimeters: [{' '.join(reflection_detections)}]",
                }
            )

        # Blind OS Command Injection
        if self.lightfuzz.interactsh_instance:
            self.lightfuzz.event_dict[self.event.url] = self.event  # Store the event associated with the URL
            for p in cmdi_probe_strings:
                _, host = self.register_interactsh_tag(
                    name="OS Command Injection",
                    description=f"OS Command Injection (OOB Interaction) Type: [{self.event.data['type']}] Parameter Name: [{self.event.data['name']}] Probe: [{p}]",
                    severity="CRITICAL",
                    confidence="CONFIRMED",
                )
                interactsh_probe = f"{p} nslookup {host} {p}"
                # we have to handle our own URL-encoding here, because our payloads include the & character
                if self.event.data["type"] == "GETPARAM":
                    interactsh_probe = urllib.parse.quote(interactsh_probe.encode(), safe="")
                # we send the probe here, and any positive detections are processed in the interactsh_callback defined in lightfuzz.py
                await self.standard_probe(
                    self.event.data["type"],
                    cookies,
                    f"{probe_value}{interactsh_probe}",
                    timeout=15,
                    skip_urlencoding=True,
                )

    async def _arith_confirm(self, http_compare, cookies, probe_str, expected, raw_marker, delim, shell_label):
        """Send a shell-arithmetic confirmation probe and return True iff the
        expected product value appears in the response and the literal
        multiplicand expression (``{a}*{b}``) does not — the latter rules out
        bulk reflection in a shell-agnostic way."""
        if self.event.data["type"] == "GETPARAM":
            probe_str = urllib.parse.quote(probe_str.encode(), safe="")
        try:
            probe = await self.compare_probe(
                http_compare,
                self.event.data["type"],
                probe_str,
                cookies,
                skip_urlencoding=True,
            )
        except HttpCompareError as e:
            self.debug(f"arithmetic probe [{shell_label}] error for [{delim}]: {e}")
            return False
        matched = probe[3] is not None and expected in probe[3].text and raw_marker not in probe[3].text
        if matched:
            self.debug(f"{shell_label} arithmetic canary [{expected}] matched for delimiter [{delim}]")
        return matched
