1"""Chapter 11 practice: a tiny object store with metadata, prefix listing, presigned URLs
2and multipart upload (Level 3 Build), and the 'signature doesn't cover the expiry' bug
3(Level 4 Debug). Standard library only. Run: python3 ch11_blob_storage.py"""
4import hashlib
5import hmac
6from urllib.parse import parse_qs, urlencode, urlparse
7
8SECRET = b"server-side-secret-key" # known only to the storage service and your backend
9
10
11class ObjectStore:
12 def __init__(self):
13 self.objects = {} # key -> {"data", "content_type", "etag", "meta"}
14 self.uploads = {} # upload_id -> {part_number: bytes}
15
16 # -- basic object API --------------------------------------------------------------
17 def put(self, key, data, content_type="application/octet-stream", meta=None):
18 etag = hashlib.md5(data).hexdigest()
19 self.objects[key] = {"data": data, "content_type": content_type,
20 "etag": etag, "meta": dict(meta or {})}
21 return etag
22
23 def get(self, key):
24 return self.objects[key]["data"]
25
26 def head(self, key):
27 o = self.objects[key]
28 return {"size": len(o["data"]), "content_type": o["content_type"],
29 "etag": o["etag"], **o["meta"]}
30
31 def list(self, prefix=""):
32 return sorted(k for k in self.objects if k.startswith(prefix))
33
34 # -- multipart upload --------------------------------------------------------------
35 def start_multipart(self, key):
36 upload_id = hashlib.sha1(f"{key}:{len(self.uploads)}".encode()).hexdigest()[:12]
37 self.uploads[upload_id] = {"key": key, "parts": {}}
38 return upload_id
39
40 def upload_part(self, upload_id, part_number, data):
41 self.uploads[upload_id]["parts"][part_number] = data
42 return hashlib.md5(data).hexdigest()
43
44 def complete_multipart(self, upload_id, content_type="application/octet-stream"):
45 up = self.uploads.pop(upload_id)
46 data = b"".join(up["parts"][n] for n in sorted(up["parts"]))
47 return self.put(up["key"], data, content_type)
48
49
50# ---- Presigned URLs -------------------------------------------------------------------------
51
52def _sign(message):
53 return hmac.new(SECRET, message.encode(), hashlib.sha256).hexdigest()
54
55
56def presign(method, key, expires_at):
57 """Your backend makes this; the client uses it to talk to storage directly."""
58 sig = _sign(f"{method}\n{key}\n{expires_at}")
59 return f"https://storage.example.com/{key}?" + urlencode({"expires": expires_at, "sig": sig})
60
61
62def verify(method, url, now):
63 """The storage service checks the signature and the expiry before allowing the request."""
64 parts = urlparse(url)
65 key = parts.path.lstrip("/")
66 q = {k: v[0] for k, v in parse_qs(parts.query).items()}
67 expected = _sign(f"{method}\n{key}\n{q['expires']}")
68 if not hmac.compare_digest(expected, q["sig"]):
69 return False
70 return now < int(q["expires"])
71
72
73# ---- Level 4: the signature that forgot the expiry ---------------------------------------------
74
75def buggy_presign(method, key, expires_at):
76 sig = _sign(f"{method}\n{key}") # BUG: expiry isn't signed
77 return f"https://storage.example.com/{key}?" + urlencode({"expires": expires_at, "sig": sig})
78
79
80def buggy_verify(method, url, now):
81 parts = urlparse(url)
82 key = parts.path.lstrip("/")
83 q = {k: v[0] for k, v in parse_qs(parts.query).items()}
84 if not hmac.compare_digest(_sign(f"{method}\n{key}"), q["sig"]):
85 return False
86 return now < int(q["expires"])
87
88
89def run_checks():
90 store = ObjectStore()
91
92 # Keys look like paths, but the namespace is flat: "folders" are just prefixes.
93 store.put("covers/42/large.jpg", b"\xff\xd8 large", "image/jpeg", {"book-id": "42"})
94 store.put("covers/42/small.jpg", b"\xff\xd8 small", "image/jpeg")
95 store.put("ebooks/42.pdf", b"%PDF-1.7 ...", "application/pdf")
96 assert store.list("covers/42/") == ["covers/42/large.jpg", "covers/42/small.jpg"]
97 h = store.head("covers/42/large.jpg")
98 assert h["content_type"] == "image/jpeg" and h["book-id"] == "42" and h["size"] == 8
99
100 # Objects are replaced whole, never edited in place.
101 store.put("ebooks/42.pdf", b"%PDF-1.7 v2")
102 assert store.get("ebooks/42.pdf") == b"%PDF-1.7 v2"
103
104 # Presigned URLs: valid until expiry, and any tampering breaks them.
105 url = presign("GET", "ebooks/42.pdf", expires_at=1_000)
106 assert verify("GET", url, now=999)
107 assert not verify("GET", url, now=1_000) # expired
108 assert not verify("PUT", url, now=10) # wrong method
109 assert not verify("GET", url.replace("ebooks/42", "ebooks/43"), now=10) # other file
110 assert not verify("GET", url.replace("expires=1000", "expires=9999999"), now=10) # extended
111
112 # Multipart: parts can arrive in any order and be retried; the result is identical.
113 original = bytes(range(256)) * 4096 # 1 MiB of test data
114 up = store.start_multipart("videos/author-talk.mp4")
115 chunks = [original[i:i + 262_144] for i in range(0, len(original), 262_144)]
116 for n in (3, 1, 4, 2): # out of order
117 store.upload_part(up, n, chunks[n - 1])
118 store.upload_part(up, 2, chunks[1]) # a retried part just replaces itself
119 store.complete_multipart(up, "video/mp4")
120 assert store.get("videos/author-talk.mp4") == original
121
122 # Level 4: with the expiry left out of the signature, a user can extend their own link.
123 bad = buggy_presign("GET", "ebooks/42.pdf", expires_at=1_000)
124 forever = bad.replace("expires=1000", "expires=99999999999")
125 assert buggy_verify("GET", forever, now=50_000_000), "the bug: an edited link still works"
126 assert not verify("GET", forever.replace("sig=", "sig=x"), now=50_000_000)
127
128 print("All checks pass.", url)
129
130
131if __name__ == "__main__":
132 run_checks()