async def _getConnection(self, user, acceptTos=False):
"""
Get a cached connection or establish one.
Can throw TermsNotAccepted via BawwabSSHClient.
"""
async with self._locks[user]:
c = self._conns.get(user, None)
if c is None:
conn = await asyncssh.connect(
client_factory=lambda: BawwabSSHClient(user, acceptTos),
host=self.host,
port=22,
username=user.name,
password=user.password,
options=asyncssh.SSHClientConnectionOptions(
known_hosts=self.knownHosts),
)
try:
stdin, stdout, stderr = await conn.open_session()
except asyncssh.misc.ChannelOpenError:
raise Exception('channel_open')
# get motd
try:
conn.bclient.motd = await asyncio.wait_for(stdout.read(
64 * 1024),
timeout=0.5)
except asyncio.TimeoutError:
pass
c = self._conns[user] = conn
return c
async def connect(self, options=None, **kwargs):
"""Open a connection to the test server"""
options = asyncssh.SSHClientConnectionOptions(options, gss_host=None)
return await asyncssh.connect(self._server_addr, self._server_port,
options=options, **kwargs)
async def do_test(self, settings):
vm_logger = logging.getLogger(self.name)
logger.info("Waiting for SSH port to open...")
await self.wait_for_ssh()
logger.info("Connecting to VM...")
client = await asyncssh.connect(
host='localhost',
port=5555,
options=asyncssh.SSHClientConnectionOptions(username='******',
password='******',
known_hosts=None),
loop=self.event_loop)
try:
sftp = await client.start_sftp_client()
try:
logger.info("Copy runtest.sh...")
async with sftp.open('runtest.sh', 'w') as fp:
await fp.write(TEST_SCRIPT.format(settings=settings))
await sftp.chmod('runtest.sh', 0o775)
if settings.source == 'local':
logger.info("Copy local.tar.gz...")
proc = subprocess.Popen(
['git', 'config', 'core.quotepath', 'off'],
cwd=ROOT_DIR)
proc.wait()
assert proc.returncode == 0
proc = subprocess.Popen([
'bash', '-c',
'tar -c -z -T<(git ls-tree --full-tree -r --name-only HEAD) -f-'
],
cwd=ROOT_DIR,
stdout=subprocess.PIPE)
async with sftp.open('local.tar.gz', 'wb') as fp:
while True:
buf = proc.stdout.read(1 << 20)
if not buf:
break
await fp.write(buf)
proc.wait()
assert proc.returncode == 0
finally:
sftp.exit()
proc = await client.create_process("./runtest.sh",
stderr=subprocess.STDOUT)
stdout_dumper = self.event_loop.create_task(
log_dumper(proc.stdout, vm_logger.info))
await proc.wait()
await stdout_dumper
assert proc.returncode == 0
finally:
client.close()
async def listen_reverse(cls, *, options=None, **kwargs):
"""Create a reverse SSH server for the tests to use"""
options = asyncssh.SSHClientConnectionOptions(
options=options, gss_host=None,
known_hosts=(['skey.pub'], [], []))
return await asyncssh.listen_reverse(port=0, family=socket.AF_INET,
options=options, **kwargs)
def test_client_options(self):
"""Test client connection options"""
with open('config', 'w') as f:
f.write('User newuser\nServerAliveInterval 1')
options = asyncssh.SSHClientConnectionOptions(username='******',
config='config')
self.assertEqual(options.username, 'user')
self.assertEqual(options.keepalive_interval, 1)
with open('config', 'w') as f:
f.write('ServerAliveInterval 2\nServerAliveCountMax 3\n')
options = asyncssh.SSHClientConnectionOptions(options, config='config')
self.assertEqual(options.keepalive_interval, 1)
self.assertEqual(options.keepalive_count_max, 3)
async def create_ssh_conn(self):
key = "%s:%d" % (self._url.address)
if key not in self.__class__.ssh_conns:
loop = asyncio.get_event_loop()
options = {
"known_hosts": None,
"host": self._url.host,
"port": self._url.port,
}
private_key_path = self._url.params.get("private_key")
if private_key_path:
if not os.path.isfile(private_key_path):
utils.logger.error(
"[%s] Private key file %s not found" %
(self.__class__.__name__, private_key_path))
return None
options["client_keys"] = [private_key_path]
if self._url.auth:
password = None
if ":" in self._url.auth:
username, password = self._url.auth.split(":", 1)
else:
username = self._url.auth
options["username"] = username
if password:
if private_key_path:
options["passphrase"] = password
else:
options["password"] = password
try:
options = asyncssh.SSHClientConnectionOptions(**options)
except (asyncssh.KeyImportError, asyncssh.KeyEncryptionError) as e:
utils.logger.error(
"[%s] Import private key %s failed: %s" %
(self.__class__.__name__, private_key_path, e))
return None
utils.logger.info(
"[%s] Create connection to ssh server %s:%d" %
(self.__class__.__name__, self._url.host, self._url.port))
ssh_conn = asyncssh.SSHClientConnection(loop, options, wait="auth")
transport = tunnel.TunnelTransport(self._tunnel, ssh_conn)
ssh_conn.connection_made(transport)
try:
await ssh_conn.wait_established()
except asyncssh.PermissionDenied as e:
utils.logger.error(
"[%s] Connect ssh server %s:%d auth failed: %s" %
(self.__class__.__name__, self._url.host, self._url.port,
e))
ssh_conn.abort()
await ssh_conn.wait_closed()
return None
self.__class__.ssh_conns[key] = ssh_conn
return self.__class__.ssh_conns[key]
async def run_client(port, host, user, passw, debug):
if debug == None:
debug = False
async with asyncssh.connect(host,
port=port,
options=asyncssh.SSHClientConnectionOptions(
known_hosts=None,
username=user,
password=passw)) as mainconn:
async with mainconn.create_session(mainsession) as initshell:
initialize(mainconn, debug)
await cossh.main_menu(mainconn, initshell, debug)
def connect(self):
""" Connect to control SSH server """
controlUrl = self.controlUrl
# remove null arguments
connectArgs = dict(
filter(lambda x: x[1], [('host', controlUrl.host),
('port', controlUrl.port),
('username', controlUrl.username)]))
knownHostsFile = '/etc/ssh/ssh_known_hosts'
if os.path.exists(knownHostsFile):
connectArgs['options'] = asyncssh.SSHClientConnectionOptions(
known_hosts=knownHostsFile)
return asyncssh.connect(**connectArgs)
def ssh_options_from_args(args, known_hosts):
kw = dict()
kw['gss_host'] = None
kw['known_hosts'] = known_hosts
if args.login is not None:
kw['username'] = args.login
if args.client_version is not None:
kw['client_version'] = args.client_version
if args.identity is not None:
kw['client_keys'] = list(args.identity)
if args.password is not None:
kw['password'] = args.password
return asyncssh.SSHClientConnectionOptions(**kw)
def connect(self, loop=(), options=None, **kwargs):
"""Open a connection to the test server"""
if loop == ():
loop = self.loop
options = asyncssh.SSHClientConnectionOptions(options, gss_host=None)
return (yield from asyncssh.connect(self._server_addr,
self._server_port,
loop=loop,
options=options,
**kwargs))
def listen_reverse(cls, *, loop=(), options=None, **kwargs):
"""Create a reverse SSH server for the tests to use"""
if loop == ():
loop = cls.loop
options = asyncssh.SSHClientConnectionOptions(
options=options, gss_host=None, known_hosts=(['skey.pub'], [], []))
return (yield from asyncssh.listen_reverse(port=0,
loop=loop,
family=socket.AF_INET,
options=options,
**kwargs))
async def test_invalid_password(self):
await self.ensure_start_server()
options = {
"known_hosts": None,
}
options = asyncssh.SSHClientConnectionOptions(**options)
with pytest.raises(asyncssh.PermissionDenied):
await asyncssh.connect(
"127.0.0.1",
self.port,
username=self.username,
password="******",
options=options,
)
self.ensure_stop_server()
async def dispatch_ssh_key(host_port, logintype, username, password_key, public_key):
host, port = host_port.split(":")
options = None
if logintype == "password":
options = asyncssh.SSHClientConnectionOptions(
known_hosts=None, username=username, password=password_key
)
elif logintype == "privatekey":
options = asyncssh.SSHClientConnectionOptions(
known_hosts=None,
client_username=username,
client_keys=[asyncssh.import_private_key(password_key)],
)
async with asyncssh.connect(host, int(port), options=options) as ssh_connection:
async with ssh_connection.start_sftp_client() as sftp_connection:
is_file_present = await sftp_connection.exists(".ssh/authorized_keys")
file_pointer = None
if is_file_present:
file_pointer = await sftp_connection.open(".ssh/authorized_keys", "a")
else:
file_pointer = await sftp_connection.open(".ssh/authorized_keys", "w")
await file_pointer.write(public_key)
await file_pointer.close()
async def _remote_connection(
self,
host: str,
addr: Optional[str] = None,
) -> "SSHClientConnection":
if not self.cons.get(host):
if not addr and host in self.mgr.inventory:
addr = self.mgr.inventory.get_addr(host)
if not addr:
raise OrchestratorError("host address is empty")
assert self.mgr.ssh_user
n = self.mgr.ssh_user + '@' + addr
logger.debug(
"Opening connection to {} with ssh options '{}'".format(
n, self.mgr._ssh_options))
asyncssh.set_log_level('DEBUG')
asyncssh.set_debug_level(3)
with self.redirect_log(host, addr):
try:
ssh_options = asyncssh.SSHClientConnectionOptions(
keepalive_interval=7, keepalive_count_max=3)
conn = await asyncssh.connect(
addr,
username=self.mgr.ssh_user,
client_keys=[self.mgr.tkey.name],
known_hosts=None,
config=[self.mgr.ssh_config_fname],
preferred_auth=['publickey'],
options=ssh_options)
except OSError:
raise
except asyncssh.Error:
raise
except Exception:
raise
self.cons[host] = conn
self.mgr.offline_hosts_remove(host)
return self.cons[host]
async def async_wd():
options = asyncssh.SSHClientConnectionOptions(
client_keys=self.key if self.key is not None else None,
password=self.password if self.key is None else None)
provider = asyncssh.connect
if self.jump:
self._tunnel = await asyncssh.connect(self.jump_host,
port=self.jump_port,
username=self.jump_user,
options=options)
provider = self._tunnel.connect_ssh
async with provider(self.host,
port=self.port,
keepalive_interval=60,
keepalive_count_max=9,
options=options) as conn:
async with conn.create_process(
'python3 -u /tmp/iris_wd.py',
input='\n'.join(
[self.path, self.pattern, self.ignore_pattern]),
stderr=asyncssh.STDOUT) as process:
self.process = process
line = False
while True:
try:
line = (await process.stdout.readline()).split('%')
path, isdir, change, mtime = line
log.debug(
f'Remote WD event: {path} {isdir} {change} {mtime}'
)
if change != 'D':
mtime = None
self.tasks.put(
File(path, mtime, self.holder, change))
except Exception as e:
# TODO: Probably here the conn and tunnel should be closed?
while line:
log.debug(line)
log.debug(e)
line = await process.stdout.readline()
break
async def test_exec_command(self):
if sys.platform == "win32":
# Ignore on windows currently
return
await self.ensure_start_server()
options = {
"known_hosts": None,
}
options = asyncssh.SSHClientConnectionOptions(**options)
async with asyncssh.connect(
"127.0.0.1",
self.port,
username=self.username,
password=self.password,
options=options,
) as conn:
message = "Hello ssh!"
result = await conn.run('echo "%s"' % message) # , check=True
# assert result.stdout.strip() == message
self.ensure_stop_server()
async def ssh_connect(self):
options = asyncssh.SSHClientConnectionOptions(
client_keys=self.key if self.key is not None else None,
password=self.password if self.key is None else None)
if self.jump:
self._tunnel = await asyncssh.connect(self.jump_host,
port=self.jump_port,
username=self.jump_user,
options=options)
self._conn = await self._tunnel.connect_ssh(self.host,
port=self.port,
username=self.user,
options=options)
else:
self._conn = await asyncssh.connect(self.host,
port=self.port,
username=self.user,
options=options)
# self._conn.connection_lost = self.connection_lost
return self._conn
async def connect(self):
utils.logger.info(
"[%s] Connect SSH server %s:%d..."
% (self.__class__.__name__, self._host, self._port)
)
options = {
"known_hosts": None,
"host": self._host,
"port": self._port,
"username": self._username,
}
options = asyncssh.SSHClientConnectionOptions(**options)
ssh_conn = await asyncssh.connect(
self._host, self._port, known_hosts=None, username=self._username
)
self._channel_writer, self._channel_reader, _ = await ssh_conn.open_session(
subsystem="tmate", env=(), send_env=(), encoding=None
)
ssh_conn.set_keepalive(300, 3)
return True
async def test_tcp_forward(self):
await self.ensure_start_server()
server = DemoTCPServer()
port = get_random_port()
server.listen(port)
options = {
"known_hosts": None,
}
options = asyncssh.SSHClientConnectionOptions(**options)
async with asyncssh.connect(
"127.0.0.1",
self.port,
username=self.username,
password=self.password,
options=options,
) as conn:
message = b"Hello ssh!"
reader, writer = await conn.open_connection("127.0.0.1", port)
writer.write(message + b"\n")
buffer = await reader.read(4096)
assert buffer.strip() == message
self.ensure_stop_server()
def get_ssh_conn_options(self):
return asyncssh.SSHClientConnectionOptions(
username=self.sftp_username,
password=self.sftp_password,
known_hosts=None,
x509_trusted_certs=None)
async def _init_ssh(self, init_boot_time=True) -> None:
await self.ssh_ready.wait()
if not self._conn:
self.ssh_ready.clear()
if self.ignore_known_hosts:
options = asyncssh.SSHClientConnectionOptions(
client_keys=self.pvtkey if self.pvtkey else None,
login_timeout=self.cmd_timeout,
password=self.password if not self.pvtkey else None,
known_hosts=None,
config=self.ssh_config_file)
else:
options = asyncssh.SSHClientConnectionOptions(
client_keys=self.pvtkey if self.pvtkey else None,
login_timeout=self.cmd_timeout,
password=self.password if not self.pvtkey else None,
config=self.ssh_config_file,
)
try:
if self.jump_host:
self.logger.info(
'Using jump host: {}, with username: {}, and port: {}'.
format(self.jump_host, self.jump_user, self.jump_port))
self._tunnel = await asyncssh.connect(
self.jump_host,
port=self.jump_port,
options=options,
username=self.jump_user)
self.logger.info(
f'Connection to jump host {self.jump_host} succeeded')
except Exception as e:
if self.sigend:
self._terminate()
return
self.logger.error(
f"ERROR: Cannot connect to jump host: {self.jump_host}, "
f" {str(e)}")
self.last_exception = e
self._conn = None
self._tunnel = None
self.ssh_ready.set()
return
try:
self._conn = await asyncssh.connect(self.address,
tunnel=self._tunnel,
username=self.username,
port=self.port,
options=options)
self.logger.info(
f"Connected to {self.address} at {time.time()}")
self.ssh_ready.set()
if init_boot_time:
await self.init_boot_time()
except Exception as e:
if self.sigend:
self._terminate()
return
self.logger.error(f"ERROR: Unable to connect, {str(e)}")
self.last_exception = e
self._conn = None
self._tunnel = None
self.ssh_ready.set()
return
ecr_repo = '812345574397.dkr.ecr.us-west-2.amazonaws.com'
s3_bucket = 'aft-experiment-logs'
cmd_template = (
'{{ /home/ubuntu/miniconda3/bin/python -m awscli ecr get-login-password --region us-west-2 ' +
'| docker login --username AWS --password-stdin {ecr_repo} ; }} && ' +
'docker pull {ecr_repo}/aft_exp:latest && ' +
'docker run --rm -it {ecr_repo}/aft_exp:latest --dataset {dataset} ' +
'--distribution {distribution} --test_fold_id {test_fold_id} --seed 1 --sampler {sampler} ' +
'--nthread {nthread} --ntrial {ntrial} --s3_bucket {s3_bucket}'
)
opts = asyncssh.SSHClientConnectionOptions(
client_keys='./bench.pem',
username='******',
password=None, # Do not permit password authentication
known_hosts=None
)
async def run_client(hostname, cmd, max_tries, logfile):
sleep_time = 5
with open(logfile, 'w') as f:
print('', file=f, end='')
for _ in range(max_tries):
try:
async with asyncssh.connect(hostname, options=opts) as conn:
print(f'[{hostname}] Conneciton successful')
async with conn.create_process(cmd, term_type='xterm-color',
stderr=asyncssh.STDOUT) as process:
async for line in process.stdout:
with open(logfile, 'a') as f:
async def _get_server_host_key(host,
port=(),
*,
tunnel=(),
family=(),
flags=0,
local_addr=None,
client_version=(),
kex_algs=(),
server_host_key_algs=(),
config=(),
options=None):
"""Retrieve an SSH server's host key
This is a coroutine which can be run to connect to an SSH server and
return the server host key presented during the SSH handshake.
A list of server host key algorithms can be provided to specify
which host key types the server is allowed to choose from. If the
key exchange is successful, the server host key sent during the
handshake is returned.
.. note:: Not all key exchange methods involve the server
presenting a host key. If something like GSS key
exchange is used without a server host key, this
method may return `None` even when the handshake
completes.
:param host:
The hostname or address to connect to
:param port: (optional)
The port number to connect to. If not specified, the default
SSH port is used.
:param tunnel: (optional)
An existing SSH client connection that this new connection should
be tunneled over. If set, a direct TCP/IP tunnel will be opened
over this connection to the requested host and port rather than
connecting directly via TCP. A string of the form
[user@]host[:port] may also be specified, in which case a
connection will first be made to that host and it will then be
used as a tunnel.
:param family: (optional)
The address family to use when creating the socket. By default,
the address family is automatically selected based on the host.
:param flags: (optional)
The flags to pass to getaddrinfo() when looking up the host address
:param local_addr: (optional)
The host and port to bind the socket to before connecting
:param client_version: (optional)
An ASCII string to advertise to the SSH server as the version of
this client, defaulting to `'AsyncSSH'` and its version number.
:param kex_algs: (optional)
A list of allowed key exchange algorithms in the SSH handshake,
taken from :ref:`key exchange algorithms <KexAlgs>`
:param server_host_key_algs: (optional)
A list of server host key algorithms to allow during the SSH
handshake, taken from :ref:`server host key algorithms
<PublicKeyAlgs>`.
:param config: (optional)
Paths to OpenSSH client configuration files to load. This
configuration will be used as a fallback to override the
defaults for settings which are not explcitly specified using
AsyncSSH's configuration options. If no paths are specified,
an attempt will be made to load the configuration from the file
:file:`.ssh/config`. If this argument is explicitly set to
`None`, no OpenSSH configuration files will be loaded. See
:ref:`SupportedClientConfigOptions` for details on what
configuration options are currently supported.
:param options: (optional)
Options to use when establishing the SSH client connection used
to retrieve the server host key. These options can be specified
either through this parameter or as direct keyword arguments to
this function.
:type host: `str`
:type port: `int`
:type tunnel: :class:`SSHClientConnection` or `str`
:type family: `socket.AF_UNSPEC`, `socket.AF_INET`, or `socket.AF_INET6`
:type flags: flags to pass to :meth:`getaddrinfo() <socket.getaddrinfo>`
:type local_addr: tuple of `str` and `int`
:type client_version: `str`
:type kex_algs: `str` or `list` of `str`
:type server_host_key_algs: `str` or `list` of `str`
:type config: `list` of `str`
:type options: :class:`SSHClientConnectionOptions`
:returns: An :class:`SSHKey` public key or `None`
"""
def conn_factory():
"""Return an SSH client connection factory"""
return asyncssh.SSHClientConnection(loop, options, wait='kex')
loop = asyncio.get_event_loop()
options = asyncssh.SSHClientConnectionOptions(
options,
config=config,
host=host,
port=port,
tunnel=tunnel,
family=family,
local_addr=local_addr,
known_hosts=None,
server_host_key_algs=server_host_key_algs,
x509_trusted_certs=None,
x509_trusted_cert_paths=None,
x509_purposes='any',
gss_host=None,
kex_algs=kex_algs,
client_version=client_version)
conn = await asyncssh.connection._connect(options.host, options.port, loop,
options.tunnel, options.family,
flags, options.local_addr,
conn_factory,
'Fetching server host key from')
server_host_key = conn.get_server_host_key()
try:
server_version = conn.get_extra_info('server_version')
except:
server_version = ''
conn.abort()
await conn.wait_closed()
return server_version, server_host_key
