def test_ssh_session_get_host(self):
session = ssh_info.SSHSession(None, None, "ip", None, None, False,
"user", None, None, None, [], False,
"Microsoft.Compute", None, None)
self.assertEqual("user@ip", session.get_host())
session = ssh_info.SSHSession("rg", "vm", None, None, None, False,
"user", None, None, None, [], False,
"Microsoft.HybridCompute", None, None)
self.assertEqual("user@vm", session.get_host())
def test_ssh_session(self, mock_abspath):
mock_abspath.side_effect = ["pub_path", "priv_path", "cert_path", "client_path", "proxy_path", "cred_path"]
expected_abspath_calls = [
mock.call("pub"),
mock.call("priv"),
mock.call("cert"),
mock.call("client/folder"),
mock.call("proxy/path"),
mock.call("cred/path")
]
session = ssh_info.SSHSession("rg", "vm", "ip", "pub", "priv", False, "user", "cert", "port", "client/folder", ['-v', '-E', 'path'], False, 'arc', 'proxy/path', 'cred/path', True)
mock_abspath.assert_has_calls(expected_abspath_calls)
self.assertEqual(session.resource_group_name, "rg")
self.assertEqual(session.vm_name, "vm")
self.assertEqual(session.ip, "ip")
self.assertEqual(session.public_key_file, "pub_path")
self.assertEqual(session.private_key_file, "priv_path")
self.assertEqual(session.use_private_ip, False)
self.assertEqual(session.local_user, "user")
self.assertEqual(session.port, "port")
self.assertEqual(session.ssh_args, ['-v', '-E', 'path'])
self.assertEqual(session.cert_file, "cert_path")
self.assertEqual(session.ssh_client_folder, "client_path")
self.assertEqual(session.ssh_proxy_folder, "proxy_path")
self.assertEqual(session.credentials_folder, "cred_path")
self.assertEqual(session.relay_info, None)
self.assertEqual(session.resource_type, "arc")
self.assertEqual(session.proxy_path, None)
self.assertEqual(session.delete_credentials, False)
self.assertEqual(session.winrdp, True)
def test_ssh_session(self, mock_abspath):
mock_abspath.side_effect = [
"pub_path", "priv_path", "cert_path", "client_path"
]
expected_abspath_calls = [
mock.call("pub"),
mock.call("priv"),
mock.call("cert"),
mock.call("client/folder")
]
session = ssh_info.SSHSession("rg", "vm", "ip", "pub", "priv", False,
"user", "cert", "port", "client/folder",
['-v', '-E', 'path'])
mock_abspath.assert_has_calls(expected_abspath_calls)
self.assertEqual(session.resource_group_name, "rg")
self.assertEqual(session.vm_name, "vm")
self.assertEqual(session.ip, "ip")
self.assertEqual(session.public_key_file, "pub_path")
self.assertEqual(session.private_key_file, "priv_path")
self.assertEqual(session.use_private_ip, False)
self.assertEqual(session.local_user, "user")
self.assertEqual(session.port, "port")
self.assertEqual(session.ssh_args, ['-v', '-E', 'path'])
self.assertEqual(session.cert_file, "cert_path")
self.assertEqual(session.ssh_client_folder, "client_path")
def test_do_ssh_op_aad_user_compute(self, mock_write_cert, mock_ssh_creds,
mock_get_mod_exp, mock_ip,
mock_check_files, mock_join,
mock_principal):
cmd = mock.Mock()
cmd.cli_ctx = mock.Mock()
cmd.cli_ctx.cloud = mock.Mock()
cmd.cli_ctx.cloud.name = "azurecloud"
op_info = ssh_info.SSHSession(None, None, "1.2.3.4", None, None, False,
None, None, None, None, None, None,
"Microsoft.Compute", None, None, False)
op_info.public_key_file = "publicfile"
op_info.private_key_file = "privatefile"
op_info.ssh_client_folder = "/client/folder"
mock_op = mock.Mock()
mock_check_files.return_value = "public", "private", False
mock_principal.return_value = ["username"]
mock_get_mod_exp.return_value = "modulus", "exponent"
profile = mock_ssh_creds.return_value
profile._adal_cache = True
profile.get_msal_token.return_value = "username", "certificate"
mock_join.return_value = "public-aadcert.pub"
custom._do_ssh_op(cmd, op_info, mock_op)
mock_check_files.assert_called_once_with("publicfile", "privatefile",
None, "/client/folder")
mock_ip.assert_not_called()
mock_get_mod_exp.assert_called_once_with("public")
mock_write_cert.assert_called_once_with("certificate",
"public-aadcert.pub")
mock_op.assert_called_once_with(op_info, False, True)
def test_start_rdp_connection(self, mock_terminate, mock_rdp, mock_wait,
mock_tunnel, mock_isopen, mock_getport):
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
"user", None, "port", None,
['arg1', 'arg2'], False,
"Microsoft.HybridCompute", None, None,
True)
op_info.public_key_file = "pub"
op_info.private_key_file = "priv"
op_info.cert_file = "cert"
op_info.ssh_client_folder = "client"
op_info.proxy_path = "proxy"
op_info.relay_info = "relay"
mock_getport.return_value = 1020
mock_isopen.return_value = True
ssh_pro = mock.Mock()
#ssh_pro.return_value.poll.return_value = None
mock_tunnel.return_value = ssh_pro, False
mock_wait.return_value = True, []
rdp_utils.start_rdp_connection(op_info, True, True)
mock_terminate.assert_called_once_with(ssh_pro, [], False)
#mock_rdp.assert_called_once_with(1020)
mock_tunnel.assert_called_once_with(op_info)
mock_wait.assert_called_once_with(ssh_pro, False)
def test_decide_resource_type_ip(self):
cmd = mock.Mock()
op_info = ssh_info.SSHSession(None, None, "ip", None, None, False,
None, None, None, None, [], False, None,
None, None, False)
self.assertEqual(custom._decide_resource_type(cmd, op_info),
"Microsoft.Compute")
def test_decide_resource_type_resourcetype_arc(self, mock_is_vm):
cmd = mock.Mock()
mock_is_vm.return_value = None, None, True
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
None, None, None, None, [], False,
"Microsoft.Compute", None, None, False)
self.assertEqual(custom._decide_resource_type(cmd, op_info),
"Microsoft.Compute")
def test_decide_resource_type_rg_vm_arc(self, mock_is_arc, mock_is_vm):
cmd = mock.Mock()
mock_is_vm.return_value = None, ResourceNotFoundError(), False
mock_is_arc.return_value = None, None, True
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
None, None, None, None, [], False, None,
None, None, False)
self.assertEqual(custom._decide_resource_type(cmd, op_info),
"Microsoft.HybridCompute")
def test_decide_resource_type_rg_vm_neither(self, mock_is_arc, mock_is_vm):
cmd = mock.Mock()
mock_is_vm.return_value = None, ResourceNotFoundError(), False
mock_is_arc.return_value = None, ResourceNotFoundError(), False
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
None, None, None, None, [], False, None,
None, None, False)
self.assertRaises(azclierror.ResourceNotFoundError,
custom._decide_resource_type, cmd, op_info)
def test_decide_resource_type_rg_vm_both(self, mock_is_arc, mock_is_vm):
cmd = mock.Mock()
mock_is_vm.return_value = None, None, True
mock_is_arc.return_value = None, None, True
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
None, None, None, None, [], False, None,
None, None, False)
self.assertRaises(azclierror.BadRequestError,
custom._decide_resource_type, cmd, op_info)
def test_ssh_session_build_args(self, mock_abspath):
mock_abspath.side_effect = [
"pub_path", "priv_path", "cert_path", "client_path"
]
session = ssh_info.SSHSession("rg", "vm", "ip", "pub", "priv", False,
"user", "cert", "port", "client/folder",
[])
self.assertEqual([
"-i", "priv_path", "-o", "CertificateFile=\"cert_path\"", "-p",
"port"
], session.build_args())
def test_ssh_session_build_args_hyvridcompute(self, mock_abspath):
mock_abspath.side_effect = [
"pub_path", "priv_path", "cert_path", "client_path"
]
session = ssh_info.SSHSession("rg", "vm", "ip", "pub", "priv", False,
"user", "cert", "port", "client/folder",
[], None, "Microsoft.HybridCompute",
None, None)
session.proxy_path = "proxy_path"
self.assertEqual([
"-o", "ProxyCommand=\"proxy_path\" -p port", "-i", "priv_path",
"-o", "CertificateFile=\"cert_path\""
], session.build_args())
def test_do_ssh_op_no_public_ip(self, mock_ip, mock_check_files):
cmd = mock.Mock()
mock_op = mock.Mock()
mock_ip.return_value = None
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
None, None, None, None, None, None,
"Microsoft.Compute", None, None, False)
self.assertRaises(azclierror.ResourceNotFoundError, custom._do_ssh_op,
cmd, op_info, mock_op)
mock_check_files.assert_not_called()
mock_ip.assert_called_once_with(cmd, "rg", "vm", False)
mock_op.assert_not_called()
def test_start_ssh_connection_arc(self, mock_system, mock_relay_str,
mock_call, mock_path, mock_copy_env,
mock_terminatecleanup):
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
"user", None, "port", None, ['arg1'],
False, "Microsoft.HybridCompute", None,
None, False)
op_info.public_key_file = "pub"
op_info.private_key_file = "priv"
op_info.cert_file = "cert"
op_info.ssh_client_folder = "client"
op_info.proxy_path = "proxy"
op_info.relay_info = "relay"
mock_system.return_value = 'Linux'
mock_call.return_value = 0
mock_relay_str.return_value = 'relay_string'
mock_copy_env.return_value = {
'var1': 'value1',
'var2': 'value2',
'var3': 'value3'
}
mock_path.return_value = 'ssh'
expected_command = [
'ssh', 'user@vm', '-o', 'ProxyCommand=\"proxy\" -p port', '-i',
'priv', '-o', 'CertificateFile=\"cert\"', 'arg1'
]
expected_env = {
'var1': 'value1',
'var2': 'value2',
'var3': 'value3',
'SSHPROXY_RELAY_INFO': 'relay_string'
}
ssh_utils.start_ssh_connection(op_info, False, False)
mock_relay_str.assert_called_once_with('relay')
mock_path.assert_called_once_with('ssh', 'client')
mock_call.assert_called_once_with(expected_command,
shell=False,
env=expected_env,
stderr=mock.ANY,
encoding='utf-8')
mock_terminatecleanup.assert_called_once_with(False, False, False,
None, 'cert', 'priv',
'pub', None, 0)
def test_do_ssh_arc_op_aad_user(self, mock_cert_exp, mock_start_ssh,
mock_write_cert, mock_ssh_creds,
mock_get_mod_exp, mock_check_files,
mock_join, mock_principal,
mock_get_relay_info, mock_get_proxy):
mock_get_proxy.return_value = '/path/to/proxy'
mock_get_relay_info.return_value = 'relay'
cmd = mock.Mock()
cmd.cli_ctx = mock.Mock()
cmd.cli_ctx.cloud = mock.Mock()
cmd.cli_ctx.cloud.name = "azurecloud"
mock_check_files.return_value = "public", "private", False
mock_principal.return_value = ["username"]
mock_get_mod_exp.return_value = "modulus", "exponent"
profile = mock_ssh_creds.return_value
profile._adal_cache = True
profile.get_msal_token.return_value = "username", "certificate"
mock_join.return_value = "public-aadcert.pub"
from datetime import timedelta
mock_cert_exp.return_value = timedelta(seconds=3600)
mock_op = mock.Mock()
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
None, None, "port", None, [], False,
"Microsoft.HybridCompute", None, None,
False)
op_info.public_key_file = "publicfile"
op_info.private_key_file = "privatefile"
op_info.ssh_client_folder = "client"
op_info.ssh_proxy_folder = "proxy"
custom._do_ssh_op(cmd, op_info, mock_op)
self.assertEqual(op_info.local_user, "username")
self.assertEqual(op_info.cert_file, "public-aadcert.pub")
mock_check_files.assert_called_once_with("publicfile", "privatefile",
None, "client")
mock_get_mod_exp.assert_called_once_with("public")
mock_write_cert.assert_called_once_with("certificate",
"public-aadcert.pub")
mock_get_proxy.assert_called_once_with('proxy')
mock_get_relay_info.assert_called_once_with(cmd, 'rg', 'vm', 3600)
mock_op.assert_called_once_with(op_info, False, True)
def test_start_ssh_connection_compute(self, mock_system, mock_copy_env,
mock_call, mock_path,
mock_terminatecleanup,
mock_startcleanup):
op_info = ssh_info.SSHSession("rg", "vm", "ip", None, None, False,
"user", None, "port", None,
['arg1', 'arg2', 'arg3'], False,
"Microsof.Compute", None, None, False)
op_info.public_key_file = "pub"
op_info.private_key_file = "priv"
op_info.cert_file = "cert"
op_info.ssh_client_folder = "client"
mock_system.return_value = 'Windows'
mock_call.return_value = 0
mock_path.return_value = 'ssh'
mock_copy_env.return_value = {
'var1': 'value1',
'var2': 'value2',
'var3': 'value3'
}
mock_startcleanup.return_value = 'log', [
'arg1', 'arg2', 'arg3', '-E', 'log', '-v'
], 'cleanup process'
expected_command = [
'ssh', 'user@ip', '-i', 'priv', '-o', 'CertificateFile=\"cert\"',
'-p', 'port', 'arg1', 'arg2', 'arg3', '-E', 'log', '-v'
]
expected_env = {'var1': 'value1', 'var2': 'value2', 'var3': 'value3'}
ssh_utils.start_ssh_connection(op_info, True, True)
mock_path.assert_called_once_with('ssh', 'client')
mock_startcleanup.assert_called_with('cert', 'priv', 'pub', False,
True, True,
['arg1', 'arg2', 'arg3'])
mock_call.assert_called_once_with(expected_command,
shell=True,
env=expected_env,
stderr=mock.ANY,
encoding='utf-8')
mock_terminatecleanup.assert_called_once_with(True, True, False,
'cleanup process',
'cert', 'priv', 'pub',
'log', 0)
def test_start_ssh_tunnel(self, mock_popen, mock_relay, mock_path,
mock_env):
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
"user", None, "port", None,
['arg1', 'arg2', '-v'], False,
"Microsoft.HybridCompute", None, None,
True)
op_info.public_key_file = "pub"
op_info.private_key_file = "priv"
op_info.cert_file = "cert"
op_info.ssh_client_folder = "client"
op_info.proxy_path = "proxy"
op_info.relay_info = "relay"
mock_env.return_value = {
'var1': 'value1',
'var2': 'value2',
'var3': 'value3'
}
mock_path.return_value = 'ssh'
mock_relay.return_value = 'relay_string'
mock_popen.return_value = 'ssh_process'
expected_command = [
'ssh', "user@vm", '-o', 'ProxyCommand=\"proxy\" -p port', '-i',
'priv', '-o', 'CertificateFile=\"cert\"', 'arg1', 'arg2', '-v'
]
expected_env = {
'var1': 'value1',
'var2': 'value2',
'var3': 'value3',
'SSHPROXY_RELAY_INFO': 'relay_string'
}
ssh_sub, print_logs = rdp_utils.start_ssh_tunnel(op_info)
self.assertEqual(ssh_sub, 'ssh_process')
self.assertEqual(print_logs, True)
mock_popen.assert_called_once_with(expected_command,
shell=True,
stderr=mock.ANY,
env=expected_env,
encoding='utf-8')
mock_relay.assert_called_once_with("relay")
mock_path.assert_called_once_with('ssh', 'client')
def test_do_ssh_op_arc_local_user(self, mock_get_cert, mock_check_keys,
mock_start_ssh, mock_get_relay_info,
mock_get_proxy):
cmd = mock.Mock()
mock_op = mock.Mock()
op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False,
"user", None, "port", None, [], False,
"Microsoft.HybridCompute", None, None)
op_info.private_key_file = "priv"
op_info.cert_file = "cert"
op_info.ssh_client_folder = "client"
op_info.ssh_proxy_folder = "proxy"
custom._do_ssh_op(cmd, op_info, mock_op)
mock_get_proxy.assert_called_once_with('proxy')
mock_get_relay_info.assert_called_once_with(cmd, 'rg', 'vm', None)
mock_op.assert_called_once_with(op_info, False, False)
mock_get_cert.assert_not_called()
mock_check_keys.assert_not_called()
def test_ssh_session_get_host(self):
session = ssh_info.SSHSession(None, None, "ip", None, None, False,
"user", None, None, None, [])
self.assertEqual("user@ip", session.get_host())
