class PeasonUploadTest(SimpleTestCase):
"""
Tests for Pearson upload_tsv
"""
def test_upload_tsv(self, connection_mock):
"""
Tests that upload uses the correct settings values
"""
upload.upload_tsv(EXAMS_SFTP_FILENAME)
ftp_mock = connection_mock.return_value.__enter__.return_value
ftp_mock.cd.assert_called_once_with(EXAMS_SFTP_UPLOAD_DIR)
ftp_mock.put.assert_called_once_with(EXAMS_SFTP_FILENAME)
@data(
SSHException(),
EOFError(),
ConnectionException('localhost', 22),
)
def test_retryable_exceptions(self, expected_exc, connection_mock):
"""
Test that if {exc_cls} is raised that it results in a RetryableSFTPException
"""
connection_mock.side_effect = expected_exc
with self.assertRaises(RetryableSFTPException) as cm:
upload.upload_tsv(EXAMS_SFTP_FILENAME)
assert isinstance(cm.exception, RetryableSFTPException)
assert cm.exception.__cause__ == expected_exc
def validate_public_key(self):
"""
Validates the host's public key against the known hosts file.
"""
if self.key:
# Log validation start.
start_log = logs.SSH_KEY_VALIDATION_START.format(ip=self.ip)
self._logger.debug(start_log)
# Read existing key information.
expected_name = self.key.get_name()
expected_value = self.key.asbytes()
# Query key information from the host.
name, value = self.query_public_key()
# Check key validity.
valid_key = name == expected_name and value == expected_value
if not valid_key:
# Log and raise exception for an invalid key.
failure_log = logs.SSH_KEY_VALIDATION_FAILURE.format(
ip=self.ip, expected_name=expected_name, name=name)
self._logger.warn(failure_log)
raise SSHException(failure_log)
# Log public key validation success.
sucess_log = logs.SSH_KEY_VALIDATION_SUCCESS.format(ip=self.ip)
self._logger.info(sucess_log)
else:
skip_log = logs.SSH_KEY_VALIDATION_SKIP.format(ip=self.ip)
self._logger.info(skip_log)
def _check_host(self):
hostkeys = paramiko.hostkeys.HostKeys()
hostkeys.load(Path('~').expanduser().joinpath('.ssh', 'known_hosts'))
if len(hostkeys.items()) == 0:
raise Exception('No Host Keys Found')
if hostkeys.lookup(self.host) is None:
raise SSHException(f'No hostkey for host {self.host} found.')
def test_exec_command_raises_on_ssh_exceptions(any_shell):
mock_channel = Mock()
any_shell.sshprocess.transport.open_session.return_value = mock_channel
any_shell.sshprocess.is_connected.return_value = False
mock_channel.exec_command.side_effect = SSHException('paramiko error')
with pytest.raises(exceptions.ConnectionError):
any_shell.exec_command(ANY_COMMAND)
def gen_keys(key="", key_path_dir=""):
"""
在KEY_DIR下创建一个 uuid命名的目录，
并且在该目录下 生产一对秘钥
:return: 返回目录名(uuid)
"""
key_basename = "key-" + uuid4().hex
if not key_path_dir:
key_path_dir = os.path.join(KEY_DIR, 'role_key', key_basename)
private_key = os.path.join(key_path_dir, 'id_rsa')
public_key = os.path.join(key_path_dir, 'id_rsa.pub')
mkdir(key_path_dir, mode=0755)
if not key:
key = RSAKey.generate(2048)
key.write_private_key_file(private_key)
else:
key_file = os.path.join(key_path_dir, 'id_rsa')
with open(key_file, 'w') as f:
f.write(key)
f.close()
with open(key_file) as f:
try:
key = RSAKey.from_private_key(f)
except SSHException, e:
shutil.rmtree(key_path_dir, ignore_errors=True)
raise SSHException(e)
def _connect(self):
host_id = self._host_cfg['host_id']
host, port = self._host_cfg[host_id, 'host']
user = self._host_cfg[host_id, 'user']
passwd = self._host_cfg[host_id, 'password']
timeout = self._host_cfg[host_id, 'timeout'] or None
known_hosts = self._host_cfg[host_id, 'known_hosts']
key_type = self._host_cfg[host_id, 'key_type']
key_file = self._host_cfg[host_id, 'key_file']
key_pass = self._host_cfg[host_id, 'key_pass']
try:
if key_type:
key = _KEY_TYPES[key_type](filename=key_file,
password=key_pass)
_logger.debug('private key: %s', key.get_name())
else:
key = None
hostname = utils.format_knownhost(host, port)
hostkeys = HostKeys(known_hosts)
transport = Transport((host, port))
transport.start_client(timeout=timeout)
hostkey = transport.get_remote_server_key()
if not hostkeys.check(hostname, hostkey):
raise SSHException('Incorrect hostkey')
if key:
transport.auth_publickey(user, key)
else:
transport.auth_password(user, passwd)
client = transport.open_sftp_client()
client.get_channel().settimeout(timeout)
_logger.debug('client for %s created', hostname)
return client
except (OSError, SSHException) as ex:
raise ConnectError(f'Connection to server "{host}:{port}"'
f' failed: {ex.args!s}')
def get_hostkey(self, host):
'''return the matching hostkey to use for verification for the host
indicated or raise an SSHException'''
kval = self.hostkeys.lookup(host) # None|{keytype: PKey}
if kval is None:
raise SSHException("No hostkey for host %s found." % host)
# return the pkey from the dict
return list(kval.values())[0]
def _download_mutil(self, host, port):
try:
sshconnector = SSHConnector(host=host, port=int(port), username=self.username,
password=self.password, pub_key_file_path=self.pub_key_file_path,
known_hosts_file_path=self.known_hosts_file_path)
sshconnector.sftp_get(self.remotefile, self.localfile)
except SSHException as e:
out, err = '', u"\n".join([e.strerror, traceback.format_exc()])
raise SSHException(err)
def ssh(self, cmd, auth_args=None):
if self.ssh_should_fail:
from paramiko import SSHException
raise SSHException("synthetic failure")
result = self.invoke(cmd, auth_args)
if isinstance(result, int):
return (result, "", "")
else:
return (0, result, "")
def test_transfer_part_ssh_exception(self, ssh_client_mock, transfer,
caplog):
"""SSH Exception occurs when connecting."""
client_mock = ssh_client_mock().__enter__()
client_mock.connect.side_effect = SSHException("Connection error")
with pytest.raises(TransferPartException):
transfer._transfer_part("dest", "0-100")
assert not client_mock.exec_command.call_count
assert ("SSH Error occurred when cURLing part: Connection error"
in caplog.messages)
def execute_command(self, command="ls -l", timeout=3000):
"""execute command
:param command:
:param timeout:
"""
stdin, stdout, stderr = self.session.exec_command(command)
if stderr:
raise SSHException(stderr)
self.sshclient.close()
return stdout
def get_missing_host_key_policy(policy):
if policy is None or policy == "ask":
return AskPolicy()
elif policy == "no" or policy == "off":
return WarningPolicy()
elif policy == "yes":
return StrictPolicy()
elif policy == "accept-new":
return AcceptNewPolicy()
else:
raise SSHException("Invalid value StrictHostKeyChecking={}".format(policy))
def testAcquireError(self):
""" Tests error handling when acquiring a TSM connection. """
_TransportMock.error = SSHException("")
self.assertRaises(PersistenceError, self._connectionPool.acquire)
_TransportMock.error = socket.error("")
self.assertRaises(PersistenceError, self._connectionPool.acquire)
_TransportMock.error = socket.gaierror("")
self.assertRaises(PersistenceError, self._connectionPool.acquire)
def test_init():
"""Testing initializing class and class helper functions"""
# Invalid TACC system specified
with pytest.raises(ValueError):
bad = TACCJobManager("foo", user=USER, psw=PW, mfa=123456)
# Invalid working directory specified, no tricky business allowed with ..
with pytest.raises(ValueError):
bad = TACCJobManager(SYSTEM,
user=USER,
psw=PW,
mfa=123456,
working_dir="../test-taccjm")
with pytest.raises(ValueError):
bad = TACCJobManager(SYSTEM,
user=USER,
psw=PW,
mfa=123456,
working_dir="test-taccjm/..")
with pytest.raises(ValueError):
bad = TACCJobManager(SYSTEM,
user=USER,
psw=PW,
mfa=123456,
working_dir="test-taccjm/../test")
# Command that should work, also test printing to stdout the output
assert JM._execute_command('echo test') == 'test\n'
# Tests command that fails due to SSH error, which we mock
with patch.object(SSHClient2FA,
'exec_command',
side_effect=SSHException('Mock ssh exception')):
with pytest.raises(SSHException):
JM._execute_command('echo test')
# Test commands that fails because of non-zero return code
with pytest.raises(TJMCommandError):
JM._execute_command('foo')
# Test making directory (remove it first)
test_dir = posixpath.join(JM.trash_dir, 'test')
try:
JM._execute_command(f"rmdir {test_dir}")
except:
pass
JM._mkdir(test_dir)
# Test making directory that will fail
with pytest.raises(TJMCommandError):
JM._mkdir(posixpath.join(JM.trash_dir, 'test/will/fail'))
def gen_keys(key=""):
"""
生成公钥 私钥
"""
output = StringIO.StringIO()
sbuffer = StringIO.StringIO()
key_content = {}
if not key:
try:
key = RSAKey.generate(2048)
key.write_private_key(output)
private_key = output.getvalue()
except IOError:
raise IOError('gen_keys: there was an error writing to the file')
except SSHException:
raise SSHException('gen_keys: the key is invalid')
else:
private_key = key
output.write(key)
try:
key = RSAKey.from_private_key(output)
except SSHException, e:
raise SSHException(e)
def gen_keys(key="", name='unknown'):
from paramiko import SSHException
from paramiko.rsakey import RSAKey
"""
生成公钥 私钥
"""
output = StringIO()
sbuffer = StringIO()
key_content = {}
if not key:
try:
key = RSAKey.generate(2048)
key.write_private_key(output)
private_key = output.getvalue()
except IOError:
raise IOError('gen_keys: there was an error writing to the file')
except SSHException:
raise SSHException('gen_keys: the key is invalid')
else:
private_key = key
output.write(key)
try:
key = RSAKey.from_private_key(output)
except SSHException as e:
raise SSHException(e)
for data in [key.get_name(),
" ",
key.get_base64(),
" %s@%s" % (name, os.uname()[1])]:
sbuffer.write(data)
public_key = sbuffer.getvalue()
key_content['public_key'] = public_key
key_content['private_key'] = private_key
return key_content
def test_process_ssh_exception_cd(self, process_zip_mock,
filtered_files_mock, os_path_exists_mock,
os_remove_mock):
"""Test that SSH exceptions bubble up"""
self.sftp.cd.side_effect = SSHException('exception')
processor = download.ArchivedResponseProcessor(self.sftp)
with self.assertRaises(RetryableSFTPException):
processor.process()
filtered_files_mock.assert_not_called()
self.sftp.remove.assert_not_called()
process_zip_mock.assert_not_called()
os_path_exists_mock.assert_not_called()
os_remove_mock.assert_not_called()
def _connect(self):
try:
client = paramiko.SSHClient()
# client.load_system_host_keys()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
if self.pub_key_file:
if plat == "win32":
client.connect(hostname=self.host, port=self.port, username=self.username,
password=self.password, key_filename=self.pub_key_file)
else:
client.connect(hostname=self.host, port=self.port, username=self.username, key_filename=self.pub_key_file)
else:
client.connect(hostname=self.host, port=self.port, username=self.username, password=self.password)
return client
except AuthenticationException as e:
raise AuthenticationException(self._error(e))
except SSHException as e:
raise SSHException(self._error(e))
def testWriteData(self):
""" Tests the behavior of the write data method. """
self._dataAdapter.exists = lambda: False
self._dataAdapter.writeData(StringIO(""))
self._dataAdapter.writeData(StringIO("Some test data to write..."))
self._dataAdapter.exists = lambda: True
self.assertRaises(PersistenceError, self._dataAdapter.writeData,
StringIO(""))
self._dataAdapter.exists = lambda: False
self._connectionMock.value = SimpleMock(error=IOError(""))
self.assertRaises(PersistenceError, self._dataAdapter.writeData,
StringIO("Some test data..."))
self._connectionMock.value = SimpleMock(error=SSHException(""))
self.assertRaises(PersistenceError, self._dataAdapter.writeData,
StringIO("Some test data..."))
def missing_host_key(self, client, hostname, key):
should_continue = input(
"No host key for {0} found in known_hosts, do you want to continue [y/n] ".format(
hostname,
),
)
if should_continue.lower() != "y":
raise SSHException(
"AskPolicy: No host key for {0} found in known_hosts".format(hostname),
)
else:
with HOST_KEYS_LOCK:
host_keys = client.get_host_keys()
host_keys.add(hostname, key.get_name(), key)
# The paramiko client saves host keys incorrectly whereas the host keys object does
# this correctly, so use that with the client filename variable.
# See: https://github.com/paramiko/paramiko/pull/1989
host_keys.save(client._host_keys_filename)
logger.warning("Added host key for {0} to known_hosts".format(hostname))
return
def testReadData(self):
""" Tests the retrieving of data. """
self._dataAdapter._sendCommand = lambda _, __: None
self._connectionMock.value = StringIO("")
self.assertEquals(self._dataAdapter.readData().read(), "")
self._connectionMock.value = StringIO("Some test data...")
self.assertEquals(self._dataAdapter.readData().read(),
"Some test data...")
self._dataAdapter.exists = lambda: False
self.assertRaises(PersistenceError, self._dataAdapter.readData)
self._dataAdapter.exists = lambda: True
self._connectionMock.value = SimpleMock(error=IOError(""))
self.assertRaises(PersistenceError, self._dataAdapter.readData)
self._connectionMock.value = SimpleMock(error=SSHException(""))
self.assertRaises(PersistenceError, self._dataAdapter.readData)
def testExist(self):
""" Tests the existence of the resource. """
self.assertTrue(self._dataAdapter.exists())
self._channelMock.stderr = "Error Code... ANS1092W kkk"
self.assertFalse(self._dataAdapter.exists())
self._channelMock.stderr = "ANS1083E No files have..."
self.assertFalse(self._dataAdapter.exists())
self._channelMock.stderr = ""
self._channelMock.stdout = "Error Code... ANS1217E kkk"
self.assertRaises(PersistenceError, self._dataAdapter.exists)
self._channelMock.stdout = ""
self._channelMock.methodNameResultMap = {
"exec_command": (None, SSHException())
} #W0201: Pylint cannot see this on commit time
self.assertRaises(PersistenceError, self._dataAdapter.exists)
def make_ssh_engine(address, modelname):
"""
creates an SSH engine object depending on host and modelname
:param address: IP or hostname as a string
:param modelname: MODELNAME
:return:
"""
creds = ssh_credentials[modelname]
assert len(creds.passwords) > 0
for password in creds.passwords:
try:
ssh_client = create_ssh_session_obj_from_hostname(address, creds.username, password)
if ssh_client.get_transport() is None or ssh_client is None:
raise SSHException()
logging.getLogger("ssh").info(ssh_client)
return SshEngine(ssh_client, address, creds.username, password, modelname)
except paramiko.AuthenticationException:
logging.getLogger("ssh").info("wrong password for %s with %s and %s will fallback to another one",
address, creds.username, password)
raise Exception("Not able to log in with " + str(creds) + " on " + address)
class _TransportMock(object):
""" Mocks the paramiko Transport class. """
error = SSHException("")
def __init__(self, _):
""" Mocks constructor. """
self.username = None
self.password = None
def connect(self, username=None, password=None):
""" Mocks connect method. """
self.username = username
self.password = password
if not self.error is None:
raise self.error
def close(self):
pass
def start_server(host=None, port=None, keyfile=None):
'''
The SFTP_HOST_KEY setting is required for configuring SFTP access.
The SFTP_PORT setting defaults to 2200.
See: tardis/default_settings/sftp.py
'''
if host is None:
current_site = Site.objects.get_current()
host = current_site.domain
port = port or getattr(settings, 'SFTP_PORT', 2200)
try:
host_key = RSAKey.from_private_key(keyfile
or StringIO(settings.SFTP_HOST_KEY))
except:
raise SSHException("SSH error: failed loading SFTP host key")
server = MyTSFTPTCPServer((host, port), host_key=host_key)
try:
server.serve_forever()
except (SystemExit, KeyboardInterrupt):
server.server_close()
def gen_keys(key="", key_path_dir=""):
"""
在KEY_DIR下创建一个 uuid命名的目录，
并且在该目录下 生产一对秘钥
:return: 返回目录名(uuid)
"""
key_basename = "key-" + uuid4().hex
if not key_path_dir:
key_path_dir = os.path.join(KEY_DIR, 'role_key', key_basename)
private_key = os.path.join(key_path_dir, 'id_rsa')
public_key = os.path.join(key_path_dir, 'id_rsa.pub')
mkdir(key_path_dir, mode=755)
if not key:
key = RSAKey.generate(2048)
key.write_private_key_file(private_key)
else:
key_file = os.path.join(key_path_dir, 'id_rsa')
with open(key_file, 'w') as f:
f.write(key)
f.close()
with open(key_file) as f:
try:
key = RSAKey.from_private_key(f)
except SSHException as e:
shutil.rmtree(key_path_dir, ignore_errors=True)
raise SSHException(e)
os.chmod(private_key, 0o644)
with open(public_key, 'w') as content_file:
for data in [
key.get_name(), " ",
key.get_base64(),
" %s@%s" % ("jumpserver", os.uname()[1])
]:
content_file.write(data)
return key_path_dir
def gen_keys_by_name(key="", key_path_dir="", name=""):
"""
在KEY_DIR下创建一个以用户名命名的目录，
并且在该目录下 生产一对秘钥
:return: 返回目录名(uuid)
"""
key_basename = name
if not key_path_dir:
key_path_dir = os.path.join(KEY_DIR, 'role_key', key_basename)
private_key = os.path.join(key_path_dir, 'id_rsa')
public_key = os.path.join(key_path_dir, 'id_rsa.pub')
config_file = os.path.join(key_path_dir, 'config')
mkdir(key_path_dir, mode=0755)
if not key:
key = RSAKey.generate(2048)
key.write_private_key_file(private_key)
config_content = [
'StrictHostKeyChecking no\n',
'UserKnownHostsFile /dev/null\n',
'GSSAPIAuthentication no\n'
]
config = file(config_file, 'w')
for c in config_content:
config.write(c)
config.close()
else:
key_file = os.path.join(key_path_dir, 'id_rsa')
with open(key_file, 'w') as f:
f.write(key)
f.close()
with open(key_file) as f:
try:
key = RSAKey.from_private_key(f)
except SSHException, e:
shutil.rmtree(key_path_dir, ignore_errors=True)
raise SSHException(e)
def missing_host_key(self, client, hostname, key):
logger.error("No host key for {0} found in known_hosts".format(hostname))
raise SSHException(
"StrictPolicy: No host key for {0} found in known_hosts".format(hostname),
)
class ArchivedResponseProcessorProcessTest(SimpleTestCase):
"""Tests around ArchivedResponseProcessor.process"""
def setUp(self):
self.sftp = MagicMock()
def test_process_success(self, process_zip_mock, filtered_files_mock,
os_path_exists_mock, os_remove_mock):
"""Test the happy path"""
processor = download.ArchivedResponseProcessor(self.sftp)
processor.process()
filtered_files_mock.assert_called_once_with()
self.sftp.remove.assert_called_once_with('a.zip')
process_zip_mock.assert_called_once_with('/tmp/a.zip')
os_path_exists_mock.assert_called_once_with('/tmp/a.zip')
os_remove_mock.assert_called_once_with('/tmp/a.zip')
def test_process_failure(self, process_zip_mock, filtered_files_mock,
os_path_exists_mock, os_remove_mock):
"""Test the unhappy path"""
process_zip_mock.return_value = False
processor = download.ArchivedResponseProcessor(self.sftp)
processor.process()
filtered_files_mock.assert_called_once_with()
self.sftp.remove.assert_not_called()
process_zip_mock.assert_called_once_with('/tmp/a.zip')
os_path_exists_mock.assert_called_once_with('/tmp/a.zip')
os_remove_mock.assert_called_once_with('/tmp/a.zip')
def test_process_exception(self, process_zip_mock, filtered_files_mock,
os_path_exists_mock, os_remove_mock):
"""Test that process() cleans up the local but not the remote on any processing exception"""
process_zip_mock.side_effect = Exception('exception')
processor = download.ArchivedResponseProcessor(self.sftp)
processor.process()
filtered_files_mock.assert_called_once_with()
self.sftp.remove.assert_not_called()
process_zip_mock.assert_called_once_with('/tmp/a.zip')
os_path_exists_mock.assert_called_once_with('/tmp/a.zip')
os_remove_mock.assert_called_once_with('/tmp/a.zip')
@ddt.data(
SSHException('exception'),
EOFError(),
)
def test_process_ssh_exception_remove(self, exc, process_zip_mock,
filtered_files_mock,
os_path_exists_mock, os_remove_mock):
"""Test that SSH exceptions bubble up"""
self.sftp.remove.side_effect = exc
processor = download.ArchivedResponseProcessor(self.sftp)
with self.assertRaises(RetryableSFTPException):
processor.process()
filtered_files_mock.assert_called_once_with()
self.sftp.remove.assert_called_once_with('a.zip')
process_zip_mock.assert_called_once_with('/tmp/a.zip')
os_path_exists_mock.assert_called_once_with('/tmp/a.zip')
os_remove_mock.assert_called_once_with('/tmp/a.zip')
def test_process_ssh_exception_cd(self, process_zip_mock,
filtered_files_mock, os_path_exists_mock,
os_remove_mock):
"""Test that SSH exceptions bubble up"""
self.sftp.cd.side_effect = SSHException('exception')
processor = download.ArchivedResponseProcessor(self.sftp)
with self.assertRaises(RetryableSFTPException):
processor.process()
filtered_files_mock.assert_not_called()
self.sftp.remove.assert_not_called()
process_zip_mock.assert_not_called()
os_path_exists_mock.assert_not_called()
os_remove_mock.assert_not_called()
def test_process_missing_local(self, process_zip_mock, filtered_files_mock,
os_path_exists_mock, os_remove_mock):
"""Test that a missing local file doesn't fail"""
os_path_exists_mock.return_value = False
processor = download.ArchivedResponseProcessor(self.sftp)
processor.process()
filtered_files_mock.assert_called_once_with()
self.sftp.remove.assert_called_once_with('a.zip')
process_zip_mock.assert_called_once_with('/tmp/a.zip')
os_path_exists_mock.assert_called_once_with('/tmp/a.zip')
os_remove_mock.assert_not_called()
def raise_exception(*args, **kwargs):
raise SSHException()
