1#!/usr/bin/env python3
2"""Capture a notebook with real OneNote in a disposable Windows clone."""
3import argparse
4import base64
5from contextlib import contextmanager
6import hashlib
7import json
8from pathlib import Path
9import signal
10import re
11import xml.etree.ElementTree as ET
12import sys
13import time
14import uuid
15import zipfile
16
17ROOT = Path(__file__).resolve().parent.parent
18sys.path.insert(0, str(ROOT / 'tools/w7'))
19import mcp_win7 as windows
20import vm
21
22
23def command(target, text, output, timeout_ms=120000):
24 result = windows.do_cmd(text, timeout_ms=timeout_ms, target=target)
25 with (output / 'commands.jsonl').open('a') as log:
26 log.write(json.dumps({'command': text, **result}) + '\n')
27 if result.get('exit') != 0 or result.get('error'):
28 raise RuntimeError(windows.text_result(result))
29
30
31def navigation_script(page, object_id=''):
32 for identifier in [page] + ([object_id] if object_id else []):
33 if not re.fullmatch(r'\{[0-9A-Fa-f-]+\}\{[0-9]+\}\{[0-9A-Fa-f]+\}', identifier):
34 raise ValueError('The native identity cannot be used for navigation.')
35 return ('OnError((exception, mode) => (FileAppend(exception.Message, "**"), ExitApp(1)))\n'
36 'DetectHiddenWindows false\napp := ComObject("OneNote.Application")\n'
37 f'app.NavigateTo("{page}", "{object_id}", false)\n'
38 'hwnd := WinWait("ahk_class Framework::CFrame ahk_exe ONENOTE.EXE",, 10)\n'
39 'if !hwnd\n throw Error("OneNote window did not appear")\n'
40 'WinMaximize(hwnd)\nWinActivate(hwnd)\n'
41 'if !WinWaitActive(hwnd,, 10)\n throw Error("OneNote window did not become active")\n'
42 f'app.NavigateTo("{page}", "{object_id}", false)\n'
43 'Sleep 300\n')
44
45
46def install_agent(name, output):
47 source = output / 'agent.py'
48 source.write_bytes((ROOT / 'tools/w7/payload/agent.py').read_bytes())
49 expected = hashlib.sha256(source.read_bytes()).hexdigest()
50 restart = output / 'restart-agent.py'
51 restart.write_bytes((ROOT / 'tools/w7/restart_agent.py').read_bytes())
52 for local, remote in [(source, r'C:\win7-agent\agent.py'),
53 (restart, r'C:\one-tests\restart-agent.py')]:
54 result = windows.do_put(local, remote, name)
55 if result.get('error'):
56 raise RuntimeError(result['error'])
57 result = windows.do_spawn(r'C:\win7-agent\vendor\python\python.exe C:\one-tests\restart-agent.py ONE-' + name.upper(), name)
58 if result.get('error'):
59 raise RuntimeError(result['error'])
60 deadline = time.monotonic() + 60
61 last_result = None
62 while time.monotonic() < deadline:
63 try:
64 health = windows.do_health(name)
65 last_result = health
66 if health.get('agent_sha256') == expected:
67 return health
68 except windows.Win7Error as error:
69 last_result = str(error)
70 time.sleep(.2)
71 (output / 'agent-startup-failure.json').write_text(json.dumps({'expected_sha256': expected, 'last_result': last_result}, indent=2))
72 raise RuntimeError('The updated clone agent did not become ready.')
73
74
75def collect_artifacts(name, output, content, partial=False):
76 command(name, 'powershell -NoProfile -Command "Compress-Archive -Force -Path C:\\one-tests\\runs\\capture\\%s -DestinationPath C:\\one-tests\\captured.zip"' % content, output)
77 archive_path = output / 'capture.zip'
78 result = windows.do_get('C:\\one-tests\\captured.zip', archive_path, name)
79 if result.get('error') or not archive_path.is_file():
80 raise RuntimeError('The native capture was not returned; inspect commands.jsonl.')
81 with zipfile.ZipFile(archive_path) as archive:
82 for entry in archive.infolist():
83 relative = Path(entry.filename.replace('\\', '/'))
84 if relative.is_absolute() or '..' in relative.parts:
85 raise ValueError('The native archive contains a path outside the capture.')
86 path = (output / 'failure-artifacts' if partial else output) / relative
87 if entry.is_dir():
88 path.mkdir(parents=True, exist_ok=True)
89 else:
90 path.parent.mkdir(parents=True, exist_ok=True)
91 path.write_bytes(archive.read(entry))
92 archive_path.unlink()
93
94
95@contextmanager
96def clone(output):
97 name = 'm6-' + uuid.uuid4().hex[:8]
98 if vm.instance_path(name).exists():
99 raise RuntimeError('The generated clone name is already in use; rerun the capture.')
100 try:
101 vm.create_instance(name)
102 (output / 'machine.json').write_text(json.dumps({'name': name, 'hostname': 'ONE-' + name.upper()}) + '\n')
103 vm.start_instance(name)
104 vm.wait_instance(name, 300)
105 (output / 'health.json').write_text(json.dumps(install_agent(name, output), indent=2) + '\n')
106 yield name
107 except BaseException:
108 if (output / 'health.json').exists():
109 try:
110 collect_artifacts(name, output, '*', partial=True)
111 except Exception as failure:
112 (output / 'artifact-failure.txt').write_text(str(failure))
113 try:
114 shot = windows.do_shot(name)
115 if shot.get('png_b64'):
116 (output / 'failure.png').write_bytes(base64.b64decode(shot['png_b64']))
117 except Exception:
118 pass
119 raise
120 finally:
121 if vm.instance_path(name).exists():
122 if vm.running(name):
123 try:
124 vm.shutdown(name, 60)
125 except (Exception, SystemExit):
126 vm.qmp(name, 'quit')
127 deadline = time.monotonic() + 10
128 while vm.running(name) and time.monotonic() < deadline:
129 time.sleep(0.1)
130 vm.delete_instance(name)
131 (output / 'teardown.json').write_text(json.dumps({'absent': not vm.instance_path(name).exists()}) + '\n')
132
133
134def capture(notebook, output, expected_pages=-1, author=None, screenshots=False, pdf=False, author_timeout=600, inspect=False, collect_notebook=False, interaction=None):
135 if inspect and interaction is not None:
136 raise ValueError('Choose manual inspection or an interaction callback.')
137 notebook = notebook.resolve(strict=True)
138 if not notebook.is_dir():
139 raise ValueError('Choose a notebook directory.')
140 if output.resolve().is_relative_to(notebook):
141 raise ValueError('Choose an output directory outside the source notebook.')
142 output.mkdir(parents=True, exist_ok=False)
143 scripts = output / 'scripts'
144 scripts.mkdir()
145 for name in ('cold.ps1', 'read.ps1'):
146 (scripts / name).write_bytes((ROOT / 'tools/native' / name).read_bytes())
147 if author is not None:
148 (scripts / 'author.ps1').write_bytes(author.read_bytes())
149 (output / 'run.json').write_text(json.dumps({
150 'notebook': str(notebook), 'expected_pages': expected_pages, 'author': str(author) if author else None,
151 'author_timeout_seconds': author_timeout,
152 'inspect': inspect,
153 'collect_notebook': collect_notebook,
154 'base': json.loads(vm.base_manifest('win7').read_text()),
155 'scripts': {name: hashlib.sha256((scripts / name).read_bytes()).hexdigest()
156 for name in sorted(p.name for p in scripts.iterdir())},
157 }, indent=2) + '\n')
158 source = []
159 transfer = output / 'transfer.zip'
160 with zipfile.ZipFile(transfer, 'w', zipfile.ZIP_DEFLATED) as archive:
161 for path in sorted(notebook.rglob('*')):
162 if not path.is_file():
163 continue
164 relative = path.relative_to(notebook).as_posix()
165 data = path.read_bytes()
166 source.append({'path': relative, 'bytes': len(data), 'sha256': hashlib.sha256(data).hexdigest(),
167 'mtime_ns': path.stat().st_mtime_ns})
168 archive.writestr(zipfile.ZipInfo.from_file(path, arcname=relative), data)
169 (output / 'source.json').write_text(json.dumps(source, indent=2) + '\n')
170 try:
171 with clone(output) as name:
172 for local, remote in [('cold.ps1', 'cold-current.ps1'), ('read.ps1', 'read-current.ps1')]:
173 result = windows.do_put(scripts / local, 'C:\\one-tests\\' + remote, name)
174 if result.get('error'):
175 raise RuntimeError(result['error'])
176 result = windows.do_put(transfer, 'C:\\one-tests\\transfer.zip', name)
177 if result.get('error'):
178 raise RuntimeError(result['error'])
179 command(name, 'powershell -NoProfile -Command "Expand-Archive -LiteralPath C:\\one-tests\\transfer.zip -DestinationPath C:\\one-tests\\runs\\capture\\notebook"', output)
180 if author is not None:
181 result = windows.do_put(scripts / 'author.ps1', 'C:\\one-tests\\author.ps1', name)
182 if result.get('error'):
183 raise RuntimeError(result['error'])
184 command(name, 'powershell -NoProfile -NonInteractive -ExecutionPolicy Bypass -File C:\\one-tests\\author.ps1 -Root C:\\one-tests\\runs\\capture -CloneHost ONE-%s' % name.upper(), output, author_timeout * 1000)
185 command(name, 'powershell -NoProfile -NonInteractive -ExecutionPolicy Bypass -File C:\\one-tests\\read-current.ps1 -Root C:\\one-tests\\runs\\capture -CloneHost ONE-%s -ExpectedPages %d%s' % (name.upper(), expected_pages, (' -UseCurrentCache' if author else '') + (' -Pdf' if pdf else '') + (' -KeepOpen' if screenshots or inspect or interaction is not None else '')), output, 600000)
186 collect_artifacts(name, output, '*' if author or collect_notebook else 'read')
187 if screenshots:
188 for page in sorted((output / 'read').glob('page-*.xml')):
189 page_id = ET.parse(page).getroot().attrib['ID']
190 result = windows.do_exec(
191 navigation_script(page_id),
192 target=name, shot_delay_ms=1500)
193 page.with_suffix('.navigation.json').write_text(json.dumps({k: v for k, v in result.items() if k != 'png_b64'}, indent=2))
194 if result.get('png_b64'):
195 page.with_suffix('.png').write_bytes(base64.b64decode(result['png_b64']))
196 if result.get('error') or result.get('exit') != 0 or not result.get('png_b64'):
197 raise RuntimeError('The native page screenshot failed; inspect its navigation record.')
198 if inspect or interaction is not None:
199 if interaction is not None:
200 interaction(name, output)
201 else:
202 print(f'Inspection ready: {name}; create {output / "finish"} to capture and close it.', flush=True)
203 deadline = time.monotonic() + 1800
204 while not (output / 'finish').exists() and time.monotonic() < deadline:
205 time.sleep(.2)
206 command(name, 'powershell -NoProfile -Command "Rename-Item C:\\one-tests\\runs\\capture\\read before-read"', output)
207 (output / 'read').rename(output / 'before-read')
208 command(name, 'powershell -NoProfile -NonInteractive -ExecutionPolicy Bypass -File C:\\one-tests\\read-current.ps1 -Root C:\\one-tests\\runs\\capture -CloneHost ONE-%s -UseCurrentCache -ExpectedPages -1%s' % (name.upper(), ' -Pdf' if pdf else ''), output, 600000)
209 collect_artifacts(name, output, '*')
210 for item in source:
211 path = notebook / item['path']
212 if path.stat().st_mtime_ns != item['mtime_ns'] or hashlib.sha256(path.read_bytes()).hexdigest() != item['sha256']:
213 raise RuntimeError('The source changed during capture; repeat from an isolated copy.')
214 finally:
215 transfer.unlink(missing_ok=True)
216
217
218if __name__ == '__main__':
219 parser = argparse.ArgumentParser(description=__doc__)
220 parser.add_argument('notebook', type=Path)
221 parser.add_argument('output', type=Path)
222 parser.add_argument('--expected-pages', type=int, default=-1)
223 parser.add_argument('--pdf', action='store_true')
224 parser.add_argument('--screenshots', action='store_true')
225 parser.add_argument('--inspect', action='store_true', help='Keep the clone open for inspection until an output/finish file appears, up to 30 minutes.')
226 parser.add_argument('--collect-notebook', action='store_true', help='Retain the VM notebook alongside its native read capture.')
227 parser.add_argument('--author', type=Path, help='Run a fixture-authoring PowerShell script in the clone before capture.')
228 parser.add_argument('--author-timeout', type=int, default=600, help='Authoring deadline in seconds.')
229 args = parser.parse_args()
230 def interrupted(_signum, _frame):
231 raise KeyboardInterrupt
232 signal.signal(signal.SIGTERM, interrupted)
233 if args.author_timeout <= 0:
234 parser.error('The authoring deadline must be positive.')
235 capture(args.notebook, args.output, args.expected_pages, args.author, args.screenshots, args.pdf, args.author_timeout, args.inspect, args.collect_notebook)
236 print('Captured native evidence in', args.output)