kangshifu-board/scripts/install-voc-database.py

"""Install only the private VOC database at the existing query skill's default path.

Default is a read-only plan. Exact repeat is a no-op; a different existing target
stops. No skill, environment, gateway, employee state or business record is changed.
"""
import argparse
import contextlib
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import sqlite3
import sys
import tempfile

sys.dont_write_bytecode = True


def digest(path):
    with Path(path).open('rb') as stream:
        result = hashlib.sha256()
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            result.update(block)
        return result.hexdigest()


def plain(path):
    path = Path(path).absolute()
    for parent in (path, *path.parents):
        if parent.exists() or parent.is_symlink():
            if parent.is_symlink() or parent.resolve() != parent:
                raise ValueError('SYMLINK_OR_REDIRECTED_PATH')
    return path


def package_manifest(package):
    entries = {}
    for line in (package / 'FILES.sha256').read_text(encoding='utf-8').splitlines():
        match = re.fullmatch(r'([a-f0-9]{64})  ([^\\]+)', line)
        if not match:
            raise ValueError('INVALID_PACKAGE_MANIFEST')
        sha, name = match.groups()
        rel = Path(name)
        if rel.is_absolute() or '..' in rel.parts or name in entries or name == 'FILES.sha256':
            raise ValueError('INVALID_PACKAGE_PATH')
        path = plain(package / rel)
        if not path.is_relative_to(package) or not path.is_file() or digest(path) != sha:
            raise ValueError('PACKAGE_HASH_MISMATCH:' + name)
        entries[name] = sha
    actual = {p.relative_to(package).as_posix() for p in package.rglob('*') if p.is_file()}
    for path in package.rglob('*'):
        plain(path)
    if actual != set(entries) | {'FILES.sha256'}:
        raise ValueError('UNEXPECTED_PACKAGE_FILES')
    release = json.loads((package / 'DATABASE.json').read_text(encoding='utf-8'))
    if release['database']['sha256'] != entries.get('library/library.sqlite'):
        raise ValueError('DATABASE_MANIFEST_MISMATCH')
    return release, entries


def verify_database(path, expected):
    if digest(path) != expected['sha256'] or path.stat().st_size != expected['bytes']:
        raise ValueError('DATABASE_HASH_MISMATCH')
    with contextlib.closing(sqlite3.connect(path.as_uri() + '?mode=ro&immutable=1', uri=True)) as con:
        con.execute('PRAGMA query_only=ON')
        meta = dict(con.execute('SELECT key,value FROM meta'))
        if meta.get('schema') != expected['schema'] or meta.get('status') != 'ready' or meta.get('generation') != expected['generation']:
            raise ValueError('DATABASE_VERSION_OR_GENERATION_MISMATCH')
        if meta.get('archive_sha256') != expected['originalArchiveSha256']:
            raise ValueError('ORIGINAL_ARCHIVE_MISMATCH')
        if con.execute('PRAGMA integrity_check').fetchone()[0] != 'ok' or con.execute('PRAGMA foreign_key_check').fetchall():
            raise ValueError('DATABASE_INTEGRITY_FAILED')
        counts = {}
        for table, count in expected['counts'].items():
            if not re.fullmatch('[a-z_]+', table):
                raise ValueError('INVALID_TABLE_NAME')
            counts[table] = con.execute('SELECT count(*) FROM ' + table).fetchone()[0]
            if counts[table] != count:
                raise ValueError('DATABASE_COUNT_MISMATCH:' + table)
        con.execute('SELECT rowid FROM retrieval_fts WHERE retrieval_fts MATCH ? LIMIT 1', ('拌',)).fetchall()
    return counts


def install(package, target_dir, apply=False):
    package, target_dir = plain(package), plain(target_dir)
    if package == target_dir or package in target_dir.parents or target_dir in package.parents:
        raise ValueError('NESTED_PACKAGE_AND_LIBRARY')
    release, entries = package_manifest(package)
    expected = release['database']
    source = package / 'library/library.sqlite'
    counts = verify_database(source, expected)
    destination = target_dir / 'library.sqlite'
    if target_dir.exists():
        if not target_dir.is_dir() or {p.name for p in target_dir.iterdir()} != {'library.sqlite'} or not destination.is_file():
            raise ValueError('TARGET_EXISTS_DIFFERENT_CONTENT')
        plain(destination)
        if digest(destination) != expected['sha256']:
            raise ValueError('TARGET_EXISTS_DIFFERENT_DATABASE')
        verify_database(destination, expected)
        operation = 'unchanged'
    else:
        operation = 'install-database'
    result = dict(ok=True, readOnly=not apply, applied=False, operation=operation,
        destination=str(destination), databaseGeneration=expected['generation'], databaseSha256=expected['sha256'],
        counts=counts, filesChecked=len(entries), databaseChecked=True, databaseWrites=0, databaseRowsChanged=0,
        skillFilesModified=False, gatewayModified=False, environmentModified=False, businessRecordsChanged=0,
        storageMode='static-snapshot', automaticUpdatesConfigured=False, productionBoardAligned=False,
        networkRequests=0, realMessageVerified=False)
    if not apply or operation == 'unchanged':
        return result
    target_dir.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    plain(target_dir.parent)
    stage = Path(tempfile.mkdtemp(prefix='.ksf-voc-database-', dir=target_dir.parent))
    os.chmod(stage, 0o700)
    try:
        target = stage / 'library.sqlite'
        shutil.copyfile(source, target)
        os.chmod(target, 0o600)
        verify_database(target, expected)
        if target_dir.exists() or target_dir.is_symlink():
            raise ValueError('CONCURRENT_TARGET_CREATED')
        stage.rename(target_dir)
        result.update(applied=True, filesystemVerified=True)
    finally:
        if stage.exists():
            shutil.rmtree(stage)
    # The staged database was verified before the directory rename; independently
    # reread the landed copy without ever opening SQLite in write mode.
    verify_database(destination, expected)
    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--package', required=True)
    parser.add_argument('--target-dir', default='/opt/data/ksf-voc-library')
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument('--check', action='store_true')
    mode.add_argument('--apply', action='store_true')
    args = parser.parse_args()
    try:
        print(json.dumps(install(args.package, args.target_dir, args.apply), ensure_ascii=False, indent=2))
    except (ValueError, OSError, KeyError, sqlite3.Error) as error:
        print(json.dumps(dict(ok=False, error=dict(code=str(error))), ensure_ascii=False), file=sys.stderr)
        raise SystemExit(1)