business-assistant/scripts/offline_release.py

"""Build and validate a single offline transfer archive, including its private dotenv."""
import argparse
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import subprocess
import tempfile
import zipfile


ROOT = Path(__file__).resolve().parents[1]
RELEASE_MANIFEST = '.offline-release.json'
SNAPSHOT_TARGET = 'output/knowledge-snapshot.zip'
TOKENIZER_TARGET = 'models/deepseek-tokenizer/tokenizer.json'


def checksum(stream):
    digest = hashlib.sha256()
    length = 0
    while block := stream.read(1024 * 1024):
        digest.update(block)
        length += len(block)
    return digest.hexdigest(), length


def safe_name(name):
    path = PurePosixPath(name)
    return bool(name and not name.startswith('/') and '\\' not in name and
                all(part not in ('', '.', '..') for part in path.parts) and
                path.as_posix() == name)


def code_files(root):
    output = subprocess.run(['git', 'ls-files', '-z', '--cached', '--others', '--exclude-standard'],
                            cwd=root, check=True, capture_output=True).stdout
    paths = []
    for raw in output.split(b'\0'):
        if not raw:
            continue
        name = os.fsdecode(raw).replace('\\', '/')
        if not safe_name(name) or name == RELEASE_MANIFEST:
            raise ValueError('invalid_project_path')
        if any(part == '.env' or part.startswith('.env.') for part in PurePosixPath(name).parts):
            continue
        path = root / name
        if path.is_file() and not path.is_symlink():
            paths.append((name, path))
    return dict(paths)


def release_files(root, snapshot):
    root = Path(root).resolve()
    snapshot = Path(snapshot).resolve()
    knowledge_manifest = json.loads((root / 'knowledge/prepared/ingestion/manifest.json').read_text(encoding='utf-8'))
    if knowledge_manifest.get('complete') is not True:
        raise ValueError('knowledge_not_published')
    files = code_files(root)
    files['__main__.py'] = root / 'scripts/offline_release_bootstrap.py'
    files['.env'] = root / '.env'
    files['config/dms_schema.json'] = root / 'config/dms_schema.json'
    files[SNAPSHOT_TARGET] = snapshot
    files[TOKENIZER_TARGET] = root / TOKENIZER_TARGET
    for relative in knowledge_manifest['model_files']:
        normalized = relative.replace('\\', '/')
        name = 'models/bge-large-zh-v1.5/' + normalized
        if not safe_name(name):
            raise ValueError('invalid_model_path')
        files[name] = root / name
    if any(not path.is_file() for path in files.values()):
        raise ValueError('missing_release_file')
    return files


def pack(root, snapshot, bundle):
    root, bundle = Path(root).resolve(), Path(bundle)
    if bundle.exists():
        raise ValueError('release_already_exists')
    files = release_files(root, snapshot)
    metadata = {}
    for name, path in files.items():
        with path.open('rb') as stream:
            digest, length = checksum(stream)
        metadata[name] = {'sha256': digest, 'bytes': length}
    bundle.parent.mkdir(parents=True, exist_ok=True)
    partial = bundle.with_name(bundle.name + '.partial')
    if partial.exists():
        raise ValueError('partial_release_exists')
    try:
        with zipfile.ZipFile(partial, 'w', allowZip64=True) as archive:
            for name, path in sorted(files.items()):
                compression = zipfile.ZIP_STORED if path.suffix in ('.bin', '.zip', '.whl') else zipfile.ZIP_DEFLATED
                archive.write(path, name, compress_type=compression)
            archive.writestr(RELEASE_MANIFEST, json.dumps({'schema': 1, 'files': metadata}, sort_keys=True))
        verify(partial)
        partial.replace(bundle)
    finally:
        partial.unlink(missing_ok=True)
    return {'files': len(files), 'bytes': sum(row['bytes'] for row in metadata.values())}


def verify(bundle):
    with zipfile.ZipFile(bundle) as archive:
        manifest_info = archive.getinfo(RELEASE_MANIFEST)
        if manifest_info.file_size > 32_000_000:
            raise ValueError('release_manifest_too_large')
        manifest = json.loads(archive.read(RELEASE_MANIFEST))
        files = manifest.get('files')
        if manifest.get('schema') != 1 or not isinstance(files, dict) or not files:
            raise ValueError('invalid_release_manifest')
        if (len(archive.infolist()) != len(files) + 1 or
                set(archive.namelist()) != set(files) | {RELEASE_MANIFEST} or
                any(not safe_name(name) or (name != '.env' and any(
                    part == '.env' or part.startswith('.env.') for part in PurePosixPath(name).parts))
                    for name in files)):
            raise ValueError('unexpected_release_members')
        for name, expected in files.items():
            info = archive.getinfo(name)
            if info.is_dir() or info.file_size != expected['bytes']:
                raise ValueError('release_size_mismatch')
            with archive.open(name) as stream:
                digest, length = checksum(stream)
            if (digest, length) != (expected['sha256'], expected['bytes']):
                raise ValueError('release_hash_mismatch')
        return {'files': len(files), 'bytes': sum(row['bytes'] for row in files.values())}


def extract(bundle, target):
    report = verify(bundle)
    target = Path(target).resolve()
    target.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(bundle) as archive:
        names = [name for name in archive.namelist() if name != RELEASE_MANIFEST]
        if any((target / name).resolve().is_relative_to(target) is False for name in names):
            raise ValueError('unsafe_release_target')
        if any((target / name).exists() for name in names):
            raise ValueError('release_target_not_empty')
        for name in names:
            destination = target / name
            destination.parent.mkdir(parents=True, exist_ok=True)
            fd, temporary = tempfile.mkstemp(prefix='.release-', dir=destination.parent)
            try:
                with os.fdopen(fd, 'wb') as output, archive.open(name) as source:
                    while block := source.read(1024 * 1024):
                        output.write(block)
                Path(temporary).replace(destination)
            finally:
                Path(temporary).unlink(missing_ok=True)
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    actions = parser.add_subparsers(dest='action', required=True)
    create = actions.add_parser('pack')
    create.add_argument('snapshot', type=Path)
    create.add_argument('bundle', type=Path)
    check = actions.add_parser('verify')
    check.add_argument('bundle', type=Path)
    unpack = actions.add_parser('extract')
    unpack.add_argument('bundle', type=Path)
    unpack.add_argument('target', type=Path)
    args = parser.parse_args()
    try:
        if args.action == 'pack':
            report = pack(ROOT, args.snapshot, args.bundle)
        elif args.action == 'verify':
            report = verify(args.bundle)
        else:
            report = extract(args.bundle, args.target)
        print(json.dumps({'ok': True, 'action': args.action, **report}))
        return 0
    except (OSError, ValueError, KeyError, TypeError, subprocess.CalledProcessError, zipfile.BadZipFile):
        print(json.dumps({'ok': False, 'action': args.action, 'code': 'offline_release_invalid'}))
        return 1


if __name__ == '__main__':
    raise SystemExit(main())