tamabox/dev/theme-builder/mock_s3.py

"""最小 S3 兼容假服务，只为调试生成器的「测试连接 / 上传素材」。

它不校验签名（签名正确性由 check_sigv4.py 与官方 SDK 对拍保证），只校验
请求形态与权限分支，所以能用来演练：桶名写错、密钥写错、前缀为空、
虚拟主机风格这些配置。

    独立跑：python dev/theme-builder/mock_s3.py --port 9900
    被复用：from mock_s3 import MockS3, BUCKET
"""
import argparse
import sys
import threading
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

BUCKET = "tb-assets"


class MockS3(BaseHTTPRequestHandler):
    store = {}
    requests = []

    def log_message(self, fmt, *args):
        pass

    # ---- helpers ----
    def _split(self):
        path = urllib.parse.unquote(self.path.split("?")[0]).lstrip("/")
        bucket, _, key = path.partition("/")
        return bucket, key

    def _key_id(self):
        # Authorization 形如
        #   AWS4-HMAC-SHA256 Credential=AKID/20261021/us-east-1/s3/aws4_request, SignedHeaders=...
        # 第一个逗号前既有算法名又有 Credential，不能按逗号切完再比前缀。
        auth = self.headers.get("Authorization") or ""
        i = auth.find("Credential=")
        if i < 0:
            return ""
        return auth[i + len("Credential="):].split("/")[0]

    def _body(self):
        n = int(self.headers.get("Content-Length") or 0)
        return self.rfile.read(n) if n else b""

    def _xml(self, status, code, message):
        body = ("<?xml version=\"1.0\"?><Error><Code>%s</Code><Message>%s</Message></Error>"
                % (code, message)).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/xml")
        if self.command != "HEAD":
            # HEAD 不能有 body：真 S3 也只是返回状态行，客户端仍会按状态码报错。
            self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

    def _check(self, bucket):
        MockS3.requests.append((self.command, self.path, dict(self.headers)))
        if self._key_id() == "BADKEY":
            self._xml(403, "SignatureDoesNotMatch",
                      "The request signature we calculated does not match")
            return False
        if bucket != BUCKET:
            self._xml(404, "NoSuchBucket", "The specified bucket does not exist")
            return False
        if not (self.headers.get("x-amz-date") or "").endswith("Z"):
            self._xml(403, "AccessDenied", "missing x-amz-date")
            return False
        return True

    # ---- verbs ----
    def do_HEAD(self):
        bucket, _ = self._split()
        if not self._check(bucket):
            return
        self.send_response(200)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_PUT(self):
        bucket, key = self._split()
        data = self._body()
        if not self._check(bucket):
            return
        MockS3.store[key] = (data, self.headers.get("Content-Type"),
                             self.headers.get("Cache-Control"))
        self.send_response(200)
        self.send_header("ETag", '"mock"')
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_GET(self):
        parsed = urllib.parse.urlsplit(self.path)
        qs = urllib.parse.parse_qs(parsed.query)
        bucket, key = self._split()
        if not key and qs.get("list-type") == ["2"]:
            if not self._check(bucket):
                return
            return self._list(qs.get("prefix", [""])[0],
                              int(qs.get("max-keys", ["1000"])[0] or 1000),
                              int(qs.get("continuation-token", ["0"])[0] or 0))
        MockS3.requests.append((self.command, self.path, dict(self.headers)))
        if key not in MockS3.store:
            self._xml(404, "NoSuchKey", "The specified key does not exist")
            return
        data, ctype, _ = MockS3.store[key]
        self.send_response(200)
        self.send_header("Content-Type", ctype or "application/octet-stream")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def _list(self, prefix, max_keys, offset):
        # 只实现 ListObjectsV2 里服务端用到的字段；分页用 offset 当 token，够测真实分页逻辑。
        keys = sorted(k for k in MockS3.store if k.startswith(prefix))
        page = keys[offset:offset + max_keys]
        truncated = offset + len(page) < len(keys)
        items = []
        for k in page:
            data = MockS3.store[k][0]
            items.append(
                "<Contents><Key>%s</Key><Size>%d</Size>"
                "<LastModified>2026-10-21T00:00:00.000Z</LastModified><ETag>&quot;mock&quot;</ETag></Contents>"
                % (k, len(data)))
        token = ("<NextContinuationToken>%d</NextContinuationToken>" % (offset + len(page))) if truncated else ""
        body = (
            '<?xml version="1.0" encoding="UTF-8"?>'
            '<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">'
            "<Name>%s</Name><Prefix>%s</Prefix><KeyCount>%d</KeyCount>"
            "<MaxKeys>%d</MaxKeys><IsTruncated>%s</IsTruncated>%s%s"
            "</ListBucketResult>"
            % (BUCKET, prefix, len(page), max_keys, "true" if truncated else "false",
               token, "".join(items))
        ).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/xml")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_DELETE(self):
        bucket, key = self._split()
        if not self._check(bucket):
            return
        MockS3.store.pop(key, None)
        self.send_response(204)
        self.send_header("Content-Length", "0")
        self.end_headers()


def spin(port=0):
    """起在后台线程，返回 (server, endpoint)。port=0 让系统挑空闲端口。"""
    srv = ThreadingHTTPServer(("127.0.0.1", port), MockS3)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    return srv, "http://127.0.0.1:%d" % srv.server_address[1]


def main():
    ap = argparse.ArgumentParser(description="TamaBox 主题生成器用的假 S3")
    ap.add_argument("--port", type=int, default=9900)
    args = ap.parse_args()
    srv, endpoint = spin(args.port)
    print("假 S3 已启动：%s" % endpoint)
    print("  桶名 %s，Access Key 随意（写 BADKEY 会返回 SignatureDoesNotMatch）" % BUCKET)
    print("Ctrl+C 停止")
    try:
        threading.Event().wait()
    except KeyboardInterrupt:
        srv.shutdown()
        print("\n已停止")


if __name__ == "__main__":
    sys.exit(main())