#!/usr/bin/env python3
"""
Solve script for DNS 0x20 exfiltration challenge.

The webshell exfiltrates data via DNS by:
  1. base64-encoding the content (with + and / replaced by -)
  2. prepending a 3-char MD5 checksum
  3. performing a DNS_A lookup for: {base64}.{checksum}.{attacker_domain}

DNS 0x20 encoding preserves the case of query labels in the response,
meaning the mixed-case payload in the DNS log IS the original base64.

However, intermediate resolvers may have corrupted some letter cases,
so we reconstruct the correct capitalisation by:
  - splitting the payload into groups of 4 base64 characters
  - trying all case variants per group
  - keeping only variants that decode to printable ASCII
  - finding the combination whose MD5[:3] matches the checksum in the log
"""

import base64
from hashlib import md5
from itertools import product

# ---------------------------------------------------------------------------
# Paste the data label and checksum from the 0x20-encoded DNS log entry here
# ---------------------------------------------------------------------------
PAYLOAD   = "Su5tE1QwdDRSBHLfBJn2zXjFAgFWcdNUzWQTAxjSLtfLodhjMtUzfq"
CHECKSUM  = "190"   # first 3 hex chars of md5(flag)
# ---------------------------------------------------------------------------


def reverse_strtr(s: str) -> str:
    """The PHP replaces + and / with -.  Reverse where possible (no-op if no dashes)."""
    # If there are no dashes the base64 had neither + nor /, nothing to do.
    # If there ARE dashes we'd need to try both — handled by the case-brute-force below.
    return s


def case_variants(group: str):
    """Yield every capitalisation variant of a base64 group that decodes to printable ASCII."""
    letter_positions = [(i, c) for i, c in enumerate(group) if c.isalpha()]
    for bits in product([True, False], repeat=len(letter_positions)):
        candidate = list(group)
        for (idx, _), upper in zip(letter_positions, bits):
            candidate[idx] = candidate[idx].upper() if upper else candidate[idx].lower()
        s = "".join(candidate)
        padding = "=" * ((4 - len(s) % 4) % 4)
        try:
            decoded = base64.b64decode(s + padding)
            if all(32 <= b < 127 for b in decoded):   # printable ASCII only
                yield s, decoded
        except Exception:
            pass


def solve(payload: str, checksum: str) -> str:
    payload = reverse_strtr(payload)

    # Split into groups of 4 (standard base64 block size)
    groups = [payload[i:i+4] for i in range(0, len(payload), 4)]

    # Collect valid byte candidates for each group
    group_candidates = []
    for i, g in enumerate(groups):
        cands = list(case_variants(g))
        if not cands:
            raise ValueError(f"No printable-ASCII candidate found for group {i} ({g!r})")
        group_candidates.append([dec for _, dec in cands])
        print(f"  group[{i:02d}] {g!r:6s} → {[dec for _, dec in cands]}")

    print(f"\nSearching {' × '.join(str(len(c)) for c in group_candidates)} combinations "
          f"for MD5[:3] == {checksum!r} …\n")

    # Enumerate combinations and check MD5 checksum
    for combo in product(*group_candidates):
        candidate_bytes = b"".join(combo)
        if md5(candidate_bytes).hexdigest()[:3] == checksum:
            return candidate_bytes.decode()

    raise ValueError("No combination matched the checksum — check PAYLOAD / CHECKSUM values.")


if __name__ == "__main__":
    print("=" * 60)
    print("DNS 0x20 Exfiltration — Solve Script")
    print("=" * 60)
    print(f"Payload  : {PAYLOAD}")
    print(f"Checksum : {CHECKSUM}\n")
    print("Candidate bytes per base64 group (printable ASCII filter):")

    flag = solve(PAYLOAD, CHECKSUM)

    print("=" * 60)
    print(f"FLAG: {flag}")
    print("=" * 60)