yaoyao/integrations/hermes-bots-bridge/file_transfer.py

import base64, hashlib, json, os, pathlib, sys, tempfile, time, uuid, fcntl, stat

os.umask(0o077)

CHUNK = 512 * 1024
HARD = 100 * 1024 * 1024
namespace = sys.argv[1]
assert len(namespace) == 64 and all(c in '0123456789abcdef' for c in namespace), 'invalid namespace'
root = pathlib.Path(tempfile.gettempdir()) / ('.yaoyao-transfers-' + namespace)
root.mkdir(mode=0o700, exist_ok=True)
assert not root.is_symlink() and root.stat().st_uid == os.getuid(), 'invalid transfer directory'
action = json.load(sys.stdin)
lock = open(root / '.lock', 'a+')
fcntl.flock(lock, fcntl.LOCK_EX)

def save(path, value):
    temporary = str(path) + '.tmp'
    with open(temporary, 'w') as stream:
        json.dump(value, stream); stream.flush(); os.fsync(stream.fileno())
    os.replace(temporary, path)

def clean(path, state):
    temporary = state.get('temporary')
    if temporary:
        try: os.unlink(temporary)
        except FileNotFoundError: pass
    path.unlink(missing_ok=True)

def sha_file(path):
    sha = hashlib.sha256()
    with open(path, 'rb') as stream:
        while True:
            data = stream.read(CHUNK)
            if not data: break
            sha.update(data)
    return sha.hexdigest()

for old in root.glob('*.json'):
    try:
        state = json.loads(old.read_text())
        if state['expiresAt'] <= time.time() or action['op'] == 'transfer-cleanup': clean(old, state)
    except (ValueError, KeyError): pass

if action['op'] == 'transfer-cleanup':
    print(json.dumps({'ok': True})); sys.exit(0)

transfer_id = action['transferId']
assert str(uuid.UUID(transfer_id)) == transfer_id, 'invalid transfer id'
state_path = root / (transfer_id + '.json')
state = json.loads(state_path.read_text()) if state_path.exists() else None
op = action['op']

if op == 'transfer-abort':
    if state: clean(state_path, state)
    result = {'ok': True}
elif op in ('transfer-read-open', 'transfer-write-open'):
    limit = action['maxBytes']
    assert type(limit) is int and 1024 * 1024 <= limit <= HARD, 'invalid transfer limit'
    signature = json.dumps(action, sort_keys=True)
    if state:
        assert state['signature'] == signature, 'transfer id conflict'
        result = state['metadata']
    else:
        assert len(list(root.glob('*.json'))) < 64, 'too many transfers'
        path = pathlib.Path(action['path']).expanduser().absolute()
        assert '\x00' not in str(path), 'invalid file path'
        if op == 'transfer-read-open':
            temporary = root / (transfer_id + '.source')
            sha = hashlib.sha256(); size = 0
            fd = os.open(path, os.O_RDONLY | os.O_NONBLOCK)
            try:
                info = os.fstat(fd)
                assert stat.S_ISREG(info.st_mode) and info.st_size <= limit, 'file exceeds transfer limit or is not a regular file'
                with os.fdopen(fd, 'rb', closefd=False) as source, open(temporary, 'xb') as target:
                    while True:
                        data = source.read(CHUNK)
                        if not data: break
                        size += len(data); assert size <= limit, 'file exceeds transfer limit'
                        sha.update(data); target.write(data)
            except BaseException:
                temporary.unlink(missing_ok=True); raise
            finally: os.close(fd)
            metadata = {'path': action['path'], 'size': size, 'sha256': sha.hexdigest()}
            state = {'kind': 'read', 'temporary': str(temporary), 'metadata': metadata}
        else:
            size = action['size']; sha = action['sha256']
            assert type(size) is int and 0 <= size <= limit, 'file exceeds transfer limit'
            assert len(sha) == 64 and all(c in '0123456789abcdef' for c in sha), 'invalid digest'
            assert type(action['overwrite']) is bool, 'invalid overwrite option'
            path.parent.mkdir(parents=True, exist_ok=True)
            path = path.parent.resolve() / path.name
            temporary = path.parent / ('.yaoyao-transfer-' + transfer_id)
            with open(temporary, 'xb'): pass
            metadata = {'path': action['path'], 'size': size, 'sha256': sha}
            state = {'kind': 'write', 'temporary': str(temporary), 'file': str(path), 'overwrite': action['overwrite'], 'received': 0, 'metadata': metadata}
        state.update(signature=signature, expiresAt=time.time() + 600)
        save(state_path, state); result = metadata
else:
    assert state, 'file transfer expired'
    metadata = state['metadata']
    if state['kind'] == 'committing':
        try:
            info = os.stat(state['file'], follow_symlinks=False)
            if info.st_ino == state['inode'] and info.st_dev == state['device'] and info.st_size == metadata['size'] and sha_file(state['file']) == metadata['sha256']:
                state['kind'] = 'complete'
        except FileNotFoundError: pass
    if op == 'transfer-status':
        result = dict(metadata, complete=state['kind'] == 'complete', received=state.get('received', 0))
    elif op == 'transfer-read':
        offset = action['offset']
        assert state['kind'] == 'read' and type(offset) is int and 0 <= offset <= metadata['size'], 'invalid read offset'
        length = min(CHUNK, metadata['size'] - offset)
        with open(state['temporary'], 'rb') as source: source.seek(offset); data = source.read(length)
        assert len(data) == length, 'incomplete source chunk'
        result = {'offset': offset, 'data': base64.b64encode(data).decode()}
    elif op == 'transfer-append':
        offset = action['offset']; encoded = action['data']
        assert state['kind'] == 'write' and type(offset) is int and offset >= 0, 'invalid write offset'
        assert type(encoded) is str and len(encoded) <= ((CHUNK + 2) // 3) * 4, 'chunk too large'
        data = base64.b64decode(encoded, validate=True)
        assert 0 < len(data) <= CHUNK and base64.b64encode(data).decode() == encoded and offset + len(data) <= metadata['size'], 'invalid chunk'
        with open(state['temporary'], 'r+b') as target:
            received = os.fstat(target.fileno()).st_size
            assert offset <= received, 'out of order chunk'
            target.seek(offset)
            if offset < received:
                assert target.read(len(data)) == data, 'conflicting repeated chunk'
            else: target.write(data); target.flush(); os.fsync(target.fileno())
            state['received'] = os.fstat(target.fileno()).st_size
        result = {'received': state['received']}
    elif op == 'transfer-finish':
        if state['kind'] != 'complete':
            assert state['kind'] == 'write' and state['received'] == metadata['size'], 'incomplete transfer'
            assert sha_file(state['temporary']) == metadata['sha256'], 'file digest mismatch'
            path = pathlib.Path(state['file'])
            assert path.parent.resolve() == path.parent, 'destination directory changed'
            info = os.stat(state['temporary']); state.update(kind='committing', inode=info.st_ino, device=info.st_dev)
            save(state_path, state)
            if state['overwrite']: os.replace(state['temporary'], path)
            else: os.link(state['temporary'], path)
            state['kind'] = 'complete'
            pathlib.Path(state['temporary']).unlink(missing_ok=True)
        result = dict(metadata, complete=True)
    else: raise ValueError('invalid transfer operation')
    state['expiresAt'] = time.time() + 600; save(state_path, state)
print(json.dumps(result))