899 views
# AsyncSSH Sample Code ###### tags: `pilot` ``` python= import asyncio, asyncssh, sys, logging, re #logging.basicConfig(level=logging.DEBUG) async def ensure_can_login_with_root(hostname, username, password, port=22): with open(f'/home/{username}/.ssh/id_rsa.pub') as f: pub_key = f.read() hostname = pub_key.split(' ')[2].rstrip() auth_dir = r'/root/.ssh' auth_path = f'{auth_dir}/authorized_keys' steps = [ (r'~\$ ', r'sudo -i'), (f'password for {username}: ', password), (r'~#', f'test -d {auth_dir} || mkdir {auth_dir}'), (r'~#', f'grep -q {hostname} {auth_path} 2>/dev/null || echo "{pub_key}" >> {auth_path}'), (r'~#', r'exit'), ] async def readuntil(stdout, separator): buf = '' while True: try: out = await asyncio.wait_for(stdout.read(1024), 5) buf += out if re.search(separator, buf): break except asyncio.TimeoutError: print(f'Expect {separator} timeout') pass async with asyncssh.connect('localhost', port=port, username=username, password=password, known_hosts=None) as conn: async with conn.create_process('bash', term_type='xterm', term_size=(80, 24)) as process: for expected, input_str in steps: await readuntil(process.stdout, expected) process.stdin.write(input_str + '\n') process.terminate() async def run(info): async with asyncssh.connect(info['hostname'], username='root', port=info['port']) as conn: result = await conn.run('id | grep -oh root') print(result.stdout) async def run_client(): info = { 'hostname': 'localhost', 'username': 'user', 'password': 'passwd', 'local_username': 'user', 'port': 22 } try: await run(info) except asyncssh.PermissionDenied as e: print(f'{e}, so try to put public key at /root/.ssh/authorized_keys') await ensure_can_login_with_root(info['hostname'], info['username'], info['password'], port=info['port']) await run(info) try: asyncio.get_event_loop().run_until_complete(run_client()) except (OSError, asyncssh.Error) as exc: sys.exit('SSH connection failed: ' + str(exc)) ```