personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4from __future__ import annotations
5
6import io
7import json
8import logging
9import socket
10import urllib.error
11import xml.etree.ElementTree as ET
12from dataclasses import dataclass
13from typing import Any
14
15from solstone.think.backup import s3_wipe
16
17ENDPOINT = "https://r2.example"
18BUCKET = "journal-backups"
19PREFIX = "users/acct/inst/"
20ACCESS_KEY = "AKIDSECRET"
21SECRET_KEY = "SECRETKEYVALUE"
22SESSION_TOKEN = "SESSIONTOKENVALUE"
23S3_NS = "http://s3.amazonaws.com/doc/2006-03-01/"
24
25
26@dataclass(frozen=True)
27class RequestRecord:
28 method: str
29 url: str
30 headers: dict[str, str]
31 body: bytes
32
33
34class FakeResponse:
35 def __init__(self, status: int, body: bytes = b"") -> None:
36 self.status = status
37 self._body = body
38
39 def __enter__(self):
40 return self
41
42 def __exit__(self, _exc_type, _exc, _tb) -> bool:
43 return False
44
45 def getcode(self) -> int:
46 return self.status
47
48 def read(self) -> bytes:
49 return self._body
50
51
52class FakeOpener:
53 def __init__(self, items: list[Any]) -> None:
54 self.items = list(items)
55 self.records: list[RequestRecord] = []
56
57 def __call__(self, request, timeout: float):
58 _ = timeout
59 self.records.append(
60 RequestRecord(
61 method=request.get_method(),
62 url=request.full_url,
63 headers={key.lower(): value for key, value in request.header_items()},
64 body=request.data or b"",
65 )
66 )
67 if not self.items:
68 raise AssertionError("unexpected S3 request")
69 item = self.items.pop(0)
70 if isinstance(item, BaseException):
71 raise item
72 return item
73
74
75def _xml(tag: str, body: str = "") -> bytes:
76 return f'<{tag} xmlns="{S3_NS}">{body}</{tag}>'.encode("utf-8")
77
78
79def _list_objects(
80 keys: list[str],
81 *,
82 truncated: bool = False,
83 token: str | None = None,
84) -> bytes:
85 contents = "".join(f"<Contents><Key>{key}</Key></Contents>" for key in keys)
86 trailer = f"<IsTruncated>{str(truncated).lower()}</IsTruncated>"
87 if token is not None:
88 trailer = f"{trailer}<NextContinuationToken>{token}</NextContinuationToken>"
89 return _xml("ListBucketResult", contents + trailer)
90
91
92def _list_uploads(
93 uploads: list[tuple[str, str]],
94 *,
95 truncated: bool = False,
96 next_key: str | None = None,
97 next_upload: str | None = None,
98) -> bytes:
99 body = "".join(
100 f"<Upload><Key>{key}</Key><UploadId>{upload_id}</UploadId></Upload>"
101 for key, upload_id in uploads
102 )
103 body = f"{body}<IsTruncated>{str(truncated).lower()}</IsTruncated>"
104 if next_key is not None:
105 body = f"{body}<NextKeyMarker>{next_key}</NextKeyMarker>"
106 if next_upload is not None:
107 body = f"{body}<NextUploadIdMarker>{next_upload}</NextUploadIdMarker>"
108 return _xml("ListMultipartUploadsResult", body)
109
110
111def _delete_result(errors: list[str] | None = None) -> bytes:
112 body = ""
113 for code in errors or []:
114 body = f"{body}<Error><Key>key</Key><Code>{code}</Code></Error>"
115 return _xml("DeleteResult", body)
116
117
118def _http_error(status: int) -> urllib.error.HTTPError:
119 return urllib.error.HTTPError(
120 f"{ENDPOINT}/{BUCKET}",
121 status,
122 "error",
123 hdrs=None,
124 fp=io.BytesIO(b""),
125 )
126
127
128def _wipe(opener: FakeOpener | Any) -> s3_wipe.WipeResult:
129 return s3_wipe.wipe_prefix(
130 endpoint=ENDPOINT,
131 bucket=BUCKET,
132 prefix=PREFIX,
133 access_key_id=ACCESS_KEY,
134 secret_access_key=SECRET_KEY,
135 session_token=SESSION_TOKEN,
136 opener=opener,
137 timeout=1,
138 )
139
140
141def _object_count(body: bytes) -> int:
142 return len(ET.fromstring(body).findall("Object"))
143
144
145def test_signed_headers_include_sts_token_on_every_request() -> None:
146 opener = FakeOpener(
147 [
148 FakeResponse(200, _list_objects([])),
149 FakeResponse(200, _list_uploads([])),
150 ]
151 )
152
153 result = _wipe(opener)
154
155 assert result == s3_wipe.WipeResult("ok", None)
156 for record in opener.records:
157 assert record.headers["x-amz-security-token"] == SESSION_TOKEN
158 assert "x-amz-security-token" in record.headers["authorization"]
159 assert record.headers["authorization"].startswith(
160 "AWS4-HMAC-SHA256 Credential="
161 )
162 assert "x-amz-date" in record.headers
163 assert "x-amz-content-sha256" in record.headers
164
165
166def test_populated_prefix_deletes_objects_then_aborts_uploads() -> None:
167 keys = [f"{PREFIX}object-{index}" for index in range(1001)]
168 opener = FakeOpener(
169 [
170 FakeResponse(200, _list_objects(keys)),
171 FakeResponse(200, _delete_result()),
172 FakeResponse(200, _delete_result()),
173 FakeResponse(200, _list_uploads([(f"{PREFIX}multipart", "upload-1")])),
174 FakeResponse(204),
175 ]
176 )
177
178 result = _wipe(opener)
179
180 assert result == s3_wipe.WipeResult("ok", None)
181 assert [record.method for record in opener.records] == [
182 "GET",
183 "POST",
184 "POST",
185 "GET",
186 "DELETE",
187 ]
188 assert _object_count(opener.records[1].body) == 1000
189 assert _object_count(opener.records[2].body) == 1
190 assert "delete=" in opener.records[1].url
191 assert "Content-md5".lower() in opener.records[1].headers
192 assert "uploadId=upload-1" in opener.records[4].url
193
194
195def test_list_objects_paginates_with_continuation_token() -> None:
196 opener = FakeOpener(
197 [
198 FakeResponse(
199 200, _list_objects([f"{PREFIX}a"], truncated=True, token="T2")
200 ),
201 FakeResponse(200, _list_objects([f"{PREFIX}b"])),
202 FakeResponse(200, _delete_result()),
203 FakeResponse(200, _list_uploads([])),
204 ]
205 )
206
207 result = _wipe(opener)
208
209 assert result.status == "ok"
210 assert [record.method for record in opener.records] == ["GET", "GET", "POST", "GET"]
211 assert "continuation-token=T2" in opener.records[1].url
212
213
214def test_multipart_upload_listing_paginates_with_markers() -> None:
215 opener = FakeOpener(
216 [
217 FakeResponse(200, _list_objects([])),
218 FakeResponse(
219 200,
220 _list_uploads(
221 [(f"{PREFIX}upload-a", "upload-a")],
222 truncated=True,
223 next_key=f"{PREFIX}next",
224 next_upload="next-upload",
225 ),
226 ),
227 FakeResponse(200, _list_uploads([(f"{PREFIX}upload-b", "upload-b")])),
228 FakeResponse(204),
229 FakeResponse(204),
230 ]
231 )
232
233 result = _wipe(opener)
234
235 assert result.status == "ok"
236 assert [record.method for record in opener.records] == [
237 "GET",
238 "GET",
239 "GET",
240 "DELETE",
241 "DELETE",
242 ]
243 assert "key-marker=users%2Facct%2Finst%2Fnext" in opener.records[2].url
244 assert "upload-id-marker=next-upload" in opener.records[2].url
245
246
247def test_empty_prefix_is_success_without_delete_or_abort() -> None:
248 opener = FakeOpener(
249 [
250 FakeResponse(200, _list_objects([])),
251 FakeResponse(200, _list_uploads([])),
252 ]
253 )
254
255 result = _wipe(opener)
256
257 assert result == s3_wipe.WipeResult("ok", None)
258 assert [record.method for record in opener.records] == ["GET", "GET"]
259
260
261def test_delete_objects_per_object_access_denied_is_failed() -> None:
262 opener = FakeOpener(
263 [
264 FakeResponse(200, _list_objects([f"{PREFIX}object"])),
265 FakeResponse(200, _delete_result(["AccessDenied"])),
266 ]
267 )
268
269 result = _wipe(opener)
270
271 assert result == s3_wipe.WipeResult("error", "failed")
272
273
274def test_delete_objects_only_no_such_key_is_idempotent_success() -> None:
275 opener = FakeOpener(
276 [
277 FakeResponse(200, _list_objects([f"{PREFIX}object"])),
278 FakeResponse(200, _delete_result(["NoSuchKey", "NoSuchKey"])),
279 FakeResponse(200, _list_uploads([])),
280 ]
281 )
282
283 result = _wipe(opener)
284
285 assert result == s3_wipe.WipeResult("ok", None)
286
287
288def test_error_mapping() -> None:
289 cases = [
290 (_http_error(403), "auth_failed"),
291 (_http_error(500), "failed"),
292 (urllib.error.URLError("down"), "unreachable"),
293 (socket.timeout("slow"), "timeout"),
294 (FakeResponse(200, b"<not-xml"), "failed"),
295 ]
296
297 for item, reason_code in cases:
298 opener = FakeOpener([item])
299
300 assert _wipe(opener) == s3_wipe.WipeResult("error", reason_code)
301
302
303def test_wipe_result_and_logs_are_secret_free(
304 caplog,
305) -> None:
306 opener = FakeOpener([_http_error(403)])
307 caplog.set_level(logging.WARNING, logger="solstone.backup.s3_wipe")
308
309 result = _wipe(opener)
310 serialized = json.dumps(result.__dict__)
311
312 for secret in (ACCESS_KEY, SECRET_KEY, SESSION_TOKEN):
313 assert secret not in serialized
314 assert secret not in repr(result)
315 assert secret not in caplog.text