attack-defense-platform/receiver/challenges/CDN.py

import io
import re
import random
import string
import logging
import subprocess
from typing import Optional

import requests
from PIL import Image

from .Challenge import Challenge


class CDN(Challenge):
    """
    SLA checker for the Flask app.

    Enforced rules (HARD FAIL if violated):
      1) /flag.txt exists in container AND equals host copy at flags/webapp.txt
      2) Image metadata generation works via ExifTool and is retrievable from /post/<id>?meta=1
      3) exiftool is installed in the container and is the provider of that metadata
    """

    # Paths & infra knobs
    flag_location = 'flags/cdn.txt'     # host copy (written by your orchestrator)
    history_location = 'history/cdn.txt'
    container_flag_path = '/flag.txt'
    container_name = 'cdn_container'            # adjust if your container name differs

    # Canonical ExifTool markers (NO app fallbacks allowed)
    _exif_must_have = (
        'File Name',
        'Date Created',
    )

    # --- Flag distribution hook (optional, mirrors your example) ---
    def distribute(self, flag: str) -> bool:
        """
        Writes/records the current flag on the host. Your infra may separately
        mount/copy it into the container; this class *also* verifies existence
        inside the container during .check().
        """
        self._make_logger()
        try:
            with open(self.flag_location, 'w') as f:
                f.write(flag)
            with open(self.history_location, 'a') as f:
                f.write(flag + '\n')
            self.logger.info(f"Flag '{flag}' written to {self.flag_location}")
            return True
        except Exception as e:
            self.logger.error(f"Failed writing host flag: {e}")
            return False

    def _make_logger(self):
        if not hasattr(self, 'logger') or self.logger is None:
            self.logger = logging.getLogger(self.__class__.__name__)
            if not self.logger.handlers:
                h = logging.StreamHandler()
                h.setFormatter(logging.Formatter('[%(levelname)s] %(message)s'))
                self.logger.addHandler(h)
                self.logger.setLevel(logging.INFO)

    # --- Helpers ---
    def _gen_username(self) -> str:
        return "user_" + ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))

    def _gen_password(self) -> str:
        return "Pw!" + ''.join(random.choices(string.ascii_letters + string.digits, k=10))

    def _tiny_png_bytes(self) -> bytes:
        img = Image.new("RGB", (2, 2), (123, 200, 50))
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()

    def _docker_exec(self, args: list[str], timeout: int = 10) -> subprocess.CompletedProcess:
        return subprocess.run(
            ["docker", "exec", self.container_name] + args,
            capture_output=True, text=True, timeout=timeout
        )

    # --- SLA core ---
    def check(self) -> bool:
        self._make_logger()
        base_url = f"http://localhost:{self.port}"
        s = requests.Session()

        # [1/9] Liveness
        self.logger.info("[1/9] Checking liveness (/login) ...")
        try:
            r = s.get(base_url + "/login", timeout=10)
            assert r.status_code == 200, f"/login HTTP {r.status_code}"
            self.logger.info("    ✓ Login page reachable")
        except Exception as e:
            self.logger.error(f"    ✗ Liveness failed: {e}")
            return False

        # [2/9] exiftool MUST exist in container
        self.logger.info("[2/9] Verifying exiftool installed in container ...")
        try:
            proc = self._docker_exec(["/bin/sh", "-lc", "command -v exiftool && exiftool -ver"])
            out = (proc.stdout or "").strip()
            assert proc.returncode == 0 and out, "exiftool not found or not runnable"
            self.logger.info(f"    ✓ exiftool OK (version: {out.splitlines()[-1]})")
        except Exception as e:
            self.logger.error(f"    ✗ exiftool check failed: {e}")
            return False

        # [3/9] flag.txt MUST exist and MUST equal host copy
        self.logger.info("[3/9] Verifying /flag.txt exists and matches host copy ...")
        try:
            proc = self._docker_exec(["/bin/sh", "-lc", f"test -f {self.container_flag_path} && cat {self.container_flag_path}"])
            assert proc.returncode == 0, "flag.txt missing in container"
            container_flag = (proc.stdout or "").strip()
            assert container_flag, "flag.txt empty"

            with open(self.flag_location, "r", encoding="utf-8") as fh:
                host_flag = fh.read().strip()

            assert container_flag == host_flag, "flag.txt content changed/mismatch"
            self.logger.info("    ✓ flag.txt present and unchanged")
        except Exception as e:
            self.logger.error(f"    ✗ Flag verification failed: {e}")
            return False

        # [4/9] Register
        self.logger.info("[4/9] Registering a fresh user ...")
        username = self._gen_username()
        password = self._gen_password()
        try:
            r = s.post(base_url + "/register",
                       data={"username": username, "password": password},
                       allow_redirects=True, timeout=10)
            assert 200 <= r.status_code < 400, f"Register HTTP {r.status_code}"
            self.logger.info(f"    ✓ Registered as {username}")
        except Exception as e:
            self.logger.error(f"    ✗ Registration failed: {e}")
            return False

        # [5/9] Login
        self.logger.info("[5/9] Logging in ...")
        try:
            r = s.post(base_url + "/login",
                       data={"username": username, "password": password},
                       allow_redirects=True, timeout=10)
            assert r.status_code == 200, f"Login HTTP {r.status_code}"
            assert "Invalid credentials" not in r.text, "Invalid credentials"
            self.logger.info("    ✓ Logged in")
        except Exception as e:
            self.logger.error(f"    ✗ Login failed: {e}")
            return False

        # [6/9] Upload tiny PNG (triggers exiftool path in app)
        self.logger.info("[6/9] Uploading PNG to generate metadata ...")
        title = "SLA Post " + ''.join(random.choices(string.ascii_letters, k=6))
        files = {"image": ("probe.png", self._tiny_png_bytes(), "image/png")}
        data = {"title": title}
        try:
            r = s.post(base_url + "/upload", files=files, data=data,
                       allow_redirects=True, timeout=20)
            assert 200 <= r.status_code < 400, f"Upload HTTP {r.status_code}"
            self.logger.info("    ✓ Upload OK")
        except Exception as e:
            self.logger.error(f"    ✗ Upload failed: {e}")
            return False

        # [7/9] Resolve post id from the JUST-CREATED post (not the newest in gallery)
        self.logger.info("[7/9] Resolving post id from upload result ...")
        try:
            created_post_id: Optional[int] = None

            # Check redirect chain and final URL for /post/<id>
            candidate_urls = []
            if r.history:
                for resp in r.history:
                    loc = resp.headers.get("Location", "")
                    if loc:
                        candidate_urls.append(loc if loc.startswith("http") else base_url + loc)
            candidate_urls.append(getattr(r, "url", ""))

            for u in candidate_urls:
                m = re.search(r'/post/(\d+)', u or "")
                if m:
                    created_post_id = int(m.group(1))
                    break

            # If still None, try to parse the response body for a /post/<id> link
            if created_post_id is None:
                mm = re.search(r'/post/(\d+)', r.text or "")
                if mm:
                    created_post_id = int(mm.group(1))

            # LAST RESORT: fallback to gallery by matching our unique title window (kept for robustness)
            if created_post_id is None:
                self.logger.info("    ↪ Falling back to /gallery title match ...")
                gr = s.get(base_url + "/gallery", timeout=10)
                if gr.status_code != 200:
                    gr = s.get(base_url + "/", timeout=10)
                    assert gr.status_code == 200, f"Index HTTP {gr.status_code}"

                post_id: Optional[int] = None
                for m in re.finditer(re.escape(title), gr.text):
                    window = gr.text[max(0, m.start()-1000): m.start()+1000]
                    mm2 = re.search(r'/post/(\d+)', window)
                    if mm2:
                        post_id = int(mm2.group(1)); break
                assert post_id is not None, "Could not resolve created post id"
                created_post_id = post_id

            self.logger.info(f"    ✓ Post id = {created_post_id}")
        except Exception as e:
            self.logger.error(f"    ✗ Post discovery failed: {e}")
            return False

        # [8/9] Post renders
        self.logger.info("[8/9] Verifying post page renders ...")
        try:
            vp = s.get(base_url + f"/post/{created_post_id}", timeout=10)
            assert vp.status_code == 200, f"Post HTTP {vp.status_code}"
            assert title in vp.text, "Post title missing"
            self.logger.info("    ✓ Post page OK")
        except Exception as e:
            self.logger.error(f"    ✗ View post failed: {e}")
            return False

        # [9/9] Metadata must appear somewhere in the post HTML (no fallback)
        self.logger.info("[9/9] Verifying metadata presence (File Name / Date Created) ...")
        try:
            mr = s.get(base_url + f"/post/{created_post_id}", timeout=10)
            assert mr.status_code == 200, f"Meta HTTP {mr.status_code}"
            html = (mr.text or "")
            assert html.strip(), "Empty metadata HTML"

            # Reject app fallbacks
            assert "no-metadata" not in html and "exif_err:" not in html, "App fallback metadata detected"

            # Required keys appear anywhere in the HTML
            assert "File Name:" in html, "Missing 'File Name:' in post HTML"
            assert "Date Created:" in html, "Missing 'Date Created:' in post HTML"

            self.logger.info("    ✓ Required metadata keys found in page HTML")
        except Exception as e:
            self.logger.error(f"    ✗ Metadata verification failed: {e}")
            return False

        self.logger.info("SLA check passed ✅")
        return True