ha-babytracker/tests/mock_server.py

"""In-process fake of the sync server, written from the app's client code.

It checks that this client speaks the protocol as the decompiled app does. It
cannot show that the real server behaves as the app expects.
"""

from __future__ import annotations

import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, Optional

COOKIE = "session=abc123"


class MockState:
    def __init__(self) -> None:
        self.account_id = 42
        self.email = "parent@example.com"
        self.password = "correct-horse"
        self.valid_sessions = True
        self.login_status: Optional[int] = None
        self.devices: list[dict[str, Any]] = []
        self.logs: dict[str, list[dict[str, Any]]] = {}
        self.uploads: list[dict[str, Any]] = []
        self.files: dict[str, bytes] = {}
        self.seed: Optional[bytes] = None
        self.requests: list[tuple[str, str]] = []
        self.login_bodies: list[dict[str, Any]] = []
        self.multipart: dict[str, tuple[str, bytes]] = {}
        # Make uploads appear in the uploader's log, as a real server would.
        self.log_uploads = False
        self.fail_with: Optional[int] = None


class Handler(BaseHTTPRequestHandler):
    server: "MockServer"

    def log_message(self, *args: Any) -> None:
        pass

    def _reply(self, status: int, body: Any = None, headers: Optional[dict[str, str]] = None) -> None:
        payload = b""
        if isinstance(body, bytes):
            payload = body
        elif body is not None:
            payload = json.dumps(body).encode("utf-8")
        self.send_response(status)
        for key, value in (headers or {}).items():
            self.send_header(key, value)
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def _body(self) -> bytes:
        return self.rfile.read(int(self.headers.get("Content-Length") or 0))

    def _authorized(self) -> bool:
        state = self.server.state
        if state.valid_sessions and COOKIE in (self.headers.get("Cookie") or ""):
            return True
        self._reply(401, b"unauthorized")
        return False

    def _login(self, body: bytes) -> None:
        state = self.server.state
        data = json.loads(body)
        state.login_bodies.append(data)
        if state.login_status:
            self._reply(state.login_status, b"nope", {"Location": "/redirected"})
            return
        if data["EmailAddress"] != state.email or data["Password"] != state.password:
            self._reply(401, b"bad credentials")
            return
        state.valid_sessions = True
        self._reply(200, {"AccountID": state.account_id}, {"Set-Cookie": f"{COOKIE}; Path=/"})

    def do_POST(self) -> None:
        state = self.server.state
        state.requests.append(("POST", self.path))
        body = self._body()
        if state.fail_with and self.path != "/session":
            self._reply(state.fail_with)
            return
        if self.path == "/session":
            self._login(body)
        elif self.path == "/redirected":
            self._reply(200, {"AccountID": 666})
        elif self.path == "/account/forgot_password":
            self._reply(200)
        elif self.path in ("/staging_account", "/account_reset"):
            state.login_bodies.append(json.loads(body))
            if json.loads(body)["EmailAddress"] == "taken@example.com":
                self._reply(409, b"exists")
            else:
                self._reply(200, None, {"Set-Cookie": f"{COOKIE}; Path=/"})
        elif not self._authorized():
            return
        elif self.path == "/device_token":
            state.multipart["token"] = ("", body)
            self._reply(201)
        elif self.path == "/account/transaction":
            data = json.loads(body)
            state.uploads.append(data)
            if state.log_uploads:
                uuid = state.login_bodies[-1]["Device"]["DeviceUUID"]
                state.logs.setdefault(uuid, []).append(data)
                state.devices = [d for d in state.devices if d["DeviceUUID"] != uuid]
                state.devices.append({"DeviceUUID": uuid, "LastSyncID": data["SyncID"]})
            self._reply(201)
        elif self.path in ("/account/picture", "/staging_account/picture", "/staging_account/seed", "/account/clonedata"):
            state.multipart[self.path] = (self.headers.get("Content-Type", ""), body)
            self._reply(201)
        else:
            self._reply(404)

    def do_PUT(self) -> None:
        state = self.server.state
        state.requests.append(("PUT", self.path))
        self._body()
        if self.path != "/staging_account":
            self._reply(404)
        elif "/staging_account/seed" not in state.multipart:
            self._reply(424, b"no seed")
        elif self._authorized():
            self._reply(200, {"AccountID": state.account_id})

    def do_DELETE(self) -> None:
        self.server.state.requests.append(("DELETE", self.path))
        if self._authorized():
            self._reply(200)

    def do_GET(self) -> None:
        state = self.server.state
        state.requests.append(("GET", self.path))
        if state.fail_with:
            self._reply(state.fail_with)
            return
        if not self._authorized():
            return
        parts = self.path.strip("/").split("/")
        if self.path == "/account/device":
            self._reply(200, state.devices)
        elif parts[:2] == ["account", "transaction"] and len(parts) == 4:
            log = state.logs.get(parts[2], [])
            after = int(parts[3])
            self._reply(200, [item for item in log if item["SyncID"] > after][:500])
        elif self.path == "/account/seed":
            if state.seed is None:
                self._reply(404, b"no seed")
            else:
                self._reply(200, state.seed, {"Content-Disposition": "attachment; filename=seed.zip"})
        elif parts[:2] == ["account", "picture"] and parts[2] in state.files:
            self._reply(200, state.files[parts[2]], {"Content-Disposition": f"attachment; filename={parts[2]}"})
        else:
            self._reply(404)


class MockServer(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self) -> None:
        super().__init__(("127.0.0.1", 0), Handler)
        self.state = MockState()
        self._thread = threading.Thread(target=self.serve_forever, args=(0.01,), daemon=True)

    @property
    def url(self) -> str:
        return f"http://127.0.0.1:{self.server_address[1]}"

    def __enter__(self) -> "MockServer":
        self._thread.start()
        return self

    def __exit__(self, *args: Any) -> None:
        self.shutdown()
        self.server_close()