"""Create scoped VDI API access through a known-host-verified SSH connection.

Requires paramiko. Password is prompted or read from VDI_PVE_SSH_PASSWORD.
The one-time API secret is written only to the requested private output file.
"""
import argparse
import getpass
import json
import os
import pathlib
import shlex
import paramiko


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--host', required=True)
    p.add_argument('--vmids', nargs='+', type=int, required=True)
    p.add_argument('--output', required=True)
    args = p.parse_args()
    target = pathlib.Path(args.output)
    if target.exists():
        raise RuntimeError('Output already exists; reuse the existing token rather than replace it')
    ssh = paramiko.SSHClient()
    ssh.load_system_host_keys(str(pathlib.Path.home() / '.ssh/known_hosts'))
    ssh.set_missing_host_key_policy(paramiko.RejectPolicy())
    password = os.environ.get('VDI_PVE_SSH_PASSWORD') or getpass.getpass('PVE SSH password: ')
    ssh.connect(args.host, username='root', password=password, look_for_keys=False, allow_agent=False, timeout=10)
    def run(*parts):
        _, out, err = ssh.exec_command(shlex.join(parts), timeout=30)
        body = out.read().decode()
        if out.channel.recv_exit_status() != 0:
            raise RuntimeError(err.read().decode())
        return body
    try:
        user, role, token = 'vdi-cloud@pve', 'VDICloudPower', 'console'
        users = json.loads(run('pveum', 'user', 'list', '--output-format', 'json'))
        if any(u['userid'] == user for u in users):
            raise RuntimeError('VDI account already exists. Inspect it and recover the existing token first')
        roles = json.loads(run('pveum', 'role', 'list', '--output-format', 'json'))
        if not any(r['roleid'] == role for r in roles):
            run('pveum', 'role', 'add', role, '--privs', 'VM.Audit VM.PowerMgmt')
        elif set(next(r for r in roles if r['roleid'] == role)['privs'].split(',')) != {'VM.Audit','VM.PowerMgmt'}:
            raise RuntimeError('Existing VDICloudPower role has unexpected privileges; inspect it first')
        run('pveum', 'user', 'add', user, '--comment', 'VDI cloud API; scoped power management')
        run('pveum', 'acl', 'modify', '/', '--users', user, '--roles', 'PVEAuditor')
        for vmid in args.vmids:
            run('pveum', 'acl', 'modify', f'/vms/{vmid}', '--users', user, '--roles', role)
        result = json.loads(run('pveum', 'user', 'token', 'add', user, token, '--privsep', '0', '--output-format', 'json'))
        target.parent.mkdir(parents=True, exist_ok=True)
        with target.open('x', encoding='utf-8') as f:
            json.dump({'PVE_URL': f'https://{args.host}:8006', 'PVE_TOKEN_ID': result['full-tokenid'], 'PVE_TOKEN_SECRET': result['value']}, f)
        os.chmod(target, 0o600)
        with ssh.open_sftp() as sftp:
            sftp.get('/etc/pve/pve-root-ca.pem', str(target.parent / 'pve-ca.pem'))
        print(f'Created dedicated account and VM power scopes {args.vmids}; credentials saved privately')
    finally:
        ssh.close()


if __name__ == '__main__':
    main()