75 lines
3.3 KiB
Python
Raw Permalink Normal View History

"""Content-addressed Runtime source and complete installed dependency manifests."""
import argparse
import hashlib
import json
import struct
from pathlib import Path
def sha(path):
return hashlib.sha256(path.read_bytes()).hexdigest()
def platform(root):
if (root / 'lib/yolo_runtime.dll').exists():
return 'windows-x64'
binary = (root / 'lib/libyolo_runtime.so').read_bytes()
if binary[:4] != b'\x7fELF':
raise ValueError('Installed Runtime is not an ELF library')
machine = struct.unpack_from('<H', binary, 18)[0]
if machine not in (183, 62):
raise ValueError('Unsupported installed Runtime architecture')
return 'linux-arm64' if machine == 183 else 'linux-x64'
def source(root):
files = [p for part in ('include', 'src', 'backends', 'adapters', 'cmake', 'scripts')
for p in (root / part).rglob('*') if p.is_file() and '__pycache__' not in p.parts]
files += [root / 'CMakeLists.txt', *root.glob('*.def')]
hashes = {p.relative_to(root).as_posix(): sha(p) for p in sorted(files)}
revision = hashlib.sha256(json.dumps(hashes, sort_keys=True).encode()).hexdigest()
return dict(schema_version=1, version='5.0.0', abi=5, source_revision=revision, files=hashes)
def verify(root, expected=None):
data = json.loads((root / 'runtime-manifest.json').read_text(encoding='utf-8'))
if data['abi'] != 5 or (expected and data['source_revision'] != expected):
raise ValueError('Runtime ABI/source revision mismatch')
for name, digest in data['files'].items():
target = (root / name).resolve()
if root.resolve() not in target.parents or not target.is_file() or sha(target) != digest:
raise ValueError('Runtime missing or changed dependency: ' + name)
required = ['include/yolo_runtime.h', 'include/yolo_runtime_types.h', 'source-identity.json']
required += (['lib/yolo_runtime.dll', 'lib/yolo_backend_onnx_cpu.dll', 'lib/onnxruntime.dll']
if data['platform'] == 'windows-x64' else ['lib/libyolo_runtime.so'])
if any(name not in data['files'] for name in required):
raise ValueError('Runtime manifest omits a required dependency')
print('RUNTIME_MANIFEST_OK', data['source_revision'])
def main():
parser = argparse.ArgumentParser()
parser.add_argument('mode', choices=['source', 'write', 'verify'])
parser.add_argument('--root', type=Path, required=True)
parser.add_argument('--output', type=Path)
parser.add_argument('--expected')
args = parser.parse_args()
root = args.root.resolve()
if args.mode == 'source':
args.output.write_text(json.dumps(source(root), indent=2), encoding='utf-8')
elif args.mode == 'write':
identity = json.loads((root / 'source-identity.json').read_text(encoding='utf-8'))
files = {p.relative_to(root).as_posix(): sha(p) for p in sorted(root.rglob('*'))
if p.is_file() and p.name != 'runtime-manifest.json'}
data = dict(schema_version=1, version=identity['version'], abi=identity['abi'],
source_revision=identity['source_revision'],
platform=platform(root), files=files)
(root / 'runtime-manifest.json').write_text(json.dumps(data, indent=2), encoding='utf-8')
verify(root)
else:
verify(root, args.expected)
if __name__ == '__main__':
main()