"""SSH helper for AD2 access. Usage: python ssh_ad2.py <cmd>"""
import subprocess, sys
import paramiko, yaml

HOST = '192.168.0.6'
USER = 'sysadmin'

def _pwd():
    r = subprocess.run(['sops','-d','D:/vault/clients/dataforth/ad2.sops.yaml'],
                       capture_output=True, text=True, timeout=30, check=True)
    return yaml.safe_load(r.stdout)['credentials']['password'].replace('\\','')

PWD = _pwd()

def run(cmd, timeout=60):
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(HOST, username=USER, password=PWD, timeout=15, look_for_keys=False, allow_agent=False)
    try:
        stdin, stdout, stderr = c.exec_command(cmd, timeout=timeout)
        out = stdout.read().decode(errors='replace')
        err = stderr.read().decode(errors='replace')
        rc = stdout.channel.recv_exit_status()
        return rc, out, err
    finally:
        c.close()

def pull(remote_path, local_path):
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(HOST, username=USER, password=PWD, timeout=15, look_for_keys=False, allow_agent=False)
    try:
        sftp = c.open_sftp()
        sftp.get(remote_path, local_path)
        sftp.close()
    finally:
        c.close()

if __name__ == '__main__':
    if sys.argv[1] == 'pull':
        pull(sys.argv[2], sys.argv[3])
        print(f'[OK] pulled {sys.argv[2]} -> {sys.argv[3]}')
    else:
        rc, out, err = run(' '.join(sys.argv[1:]))
        if out: print(out, end='')
        if err: print('STDERR:', err, file=sys.stderr, end='')
        sys.exit(rc)