File size: 1,823 Bytes
7bcd331 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | #!/usr/bin/env python3
"""Restore an archived training source tree into a new directory."""
import argparse
import hashlib
import json
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--list', action='store_true')
parser.add_argument('--snapshot')
parser.add_argument('--output', type=Path)
args = parser.parse_args()
index = json.loads((ROOT/'training/SNAPSHOTS.json').read_text())
if args.list:
print(json.dumps(index, indent=2)); return
if args.snapshot not in index or args.output is None:
parser.error('Choose --snapshot from --list and supply a new --output directory.')
output = args.output.resolve()
if output.exists():
parser.error('Output already exists; choose a new directory. No files were changed.')
manifest = json.loads((ROOT/index[args.snapshot]['snapshot']).read_text())
checked = []
for name, record in manifest['files'].items():
destination = (output/name).resolve()
source = (ROOT/record['file']).resolve()
if not destination.is_relative_to(output) or not source.is_relative_to(ROOT):
raise ValueError('Invalid snapshot path')
data = source.read_bytes()
if hashlib.sha256(data).hexdigest() != record['packaged_sha256']:
raise ValueError(f'Snapshot integrity mismatch: {name}')
checked.append((destination, data))
output.mkdir(parents=True, exist_ok=False)
for destination, data in checked:
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(data)
print(json.dumps({'snapshot': args.snapshot, 'files': len(checked), 'output': str(output)}, indent=2))
if __name__ == '__main__':
main()
|