#!/usr/bin/env python3
"""Send only S3 keys from local OpenBao to remote runtime storage over SSH.

Never forwards the OpenBao token. Runtime credentials disappear at reboot.
The existing bootstrap fallback requires --allow-bootstrap explicitly.
"""
import argparse
import json
import os
from pathlib import Path
import shlex
import subprocess
import urllib.error
import urllib.request


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--host', default='railiance01')
    ap.add_argument('--allow-bootstrap', action='store_true')
    args = ap.parse_args()
    token = os.environ.get('OPENBAO_TOKEN') or os.environ.get('VAULT_TOKEN')
    if not token:
        token = Path.home().joinpath('.vault-token').read_text().strip()
    addr = os.environ.get('BAO_ADDR', 'https://bao.coulomb.social')
    paths = ['workloads/railiance/freedom-intelligence/object-storage']
    if args.allow_bootstrap:
        paths.append('workloads/railiance/scaleway/bootstrap')
    credentials = None
    for path in paths:
        request = urllib.request.Request(f'{addr}/v1/platform/data/{path}',
                                         headers={'X-Vault-Token': token})
        try:
            with urllib.request.urlopen(request, timeout=20) as response:
                data = json.load(response)['data']['data']
        except urllib.error.HTTPError as exc:
            if exc.code == 404:
                continue
            raise RuntimeError(f'OpenBao read failed: HTTP {exc.code}') from None
        access = data.get('ACCESS_KEY') or data.get('access_key') or data.get('AWS_ACCESS_KEY_ID')
        secret = data.get('SECRET_KEY') or data.get('secret_key') or data.get('AWS_SECRET_ACCESS_KEY')
        if not access or not secret:
            raise RuntimeError('OpenBao S3 secret has unsupported fields')
        credentials = {'ACCESS_KEY': access, 'SECRET_KEY': secret}
        print(f'credential source: platform/{path}')
        break
    if credentials is None:
        raise RuntimeError('no S3 credential found')
    program = '''import json, os, pathlib, sys
data = json.load(sys.stdin)
root = pathlib.Path('/run/user') / str(os.getuid()) / 'fi-reserve'
root.mkdir(mode=0o700, exist_ok=True)
root.chmod(0o700)
fd = os.open(root / 's3.json', os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW, 0o600)
os.fchmod(fd, 0o600)
with os.fdopen(fd, 'w') as f:
    json.dump(data, f)
print('S3 credentials installed in remote runtime directory (0600)')
'''
    subprocess.run(['ssh', '-o', 'BatchMode=yes', args.host,
                    'python3 -c ' + shlex.quote(program)],
                   input=json.dumps(credentials), text=True, check=True)


if __name__ == '__main__':
    main()