diff --git a/reviewboard/scmtools/tests.py b/reviewboard/scmtools/tests.py index 4d8b3ca797ad4452f17833cd582454bb9b095168..8d825ffc0e37277c2141ed0f701faf0a2ba0b8d5 100644 --- a/reviewboard/scmtools/tests.py +++ b/reviewboard/scmtools/tests.py @@ -54,7 +54,7 @@ class SCMTestCase(SSHTestCase): raise nose.SkipTest( "Cannot perform SSH access tests. The local user's SSH " "public key must be in the %s file and SSH must be enabled." - % os.path.join(self.ssh_client.get_ssh_dir(), + % os.path.join(self.ssh_client.storage.get_ssh_dir(), 'authorized_keys')) def _test_ssh(self, repo_path, filename=None): @@ -92,7 +92,7 @@ class SCMTestCase(SSHTestCase): sshdir = os.path.join(self.tempdir, '.ssh') self._set_home(self.tempdir) - self.assertEqual(sshdir, self.ssh_client.get_ssh_dir()) + self.assertEqual(sshdir, self.ssh_client.storage.get_ssh_dir()) self.assertFalse(os.path.exists(os.path.join(sshdir, 'id_rsa'))) self.assertFalse(os.path.exists(os.path.join(sshdir, 'id_dsa'))) self.assertEqual(self.ssh_client.get_user_key(), None) @@ -118,7 +118,7 @@ class SCMTestCase(SSHTestCase): tool = repo.get_scmtool() ssh_client = SSHClient(namespace=local_site_name) - self.assertEqual(ssh_client.get_ssh_dir(), + self.assertEqual(ssh_client.storage.get_ssh_dir(), os.path.join(sshdir, local_site_name)) ssh_client.import_user_key(user_key) self.assertEqual(ssh_client.get_user_key(), user_key) diff --git a/reviewboard/ssh/client.py b/reviewboard/ssh/client.py index d4bf08a0547126519117f3c39d526da5ac3fc928..f44a7b19d298ee5de85ca221edceee179fc71704 100644 --- a/reviewboard/ssh/client.py +++ b/reviewboard/ssh/client.py @@ -1,58 +1,102 @@ import logging -import os - -from django.utils.translation import ugettext_lazy as _ +try: + from cStringIO import StringIO +except ImportError: + from StringIO import StringIO + +from django.conf import settings +from django.core.exceptions import ImproperlyConfigured +from paramiko.hostkeys import HostKeyEntry import paramiko -from reviewboard.ssh.errors import MakeSSHDirError, UnsupportedSSHKeyError +from reviewboard.ssh.errors import UnsupportedSSHKeyError -class SSHClient(paramiko.SSHClient): - _ssh_dir = None +class SSHHostKeys(paramiko.HostKeys): + """Manages known lists of host keys. - def __init__(self, namespace=None): - super(SSHClient, self).__init__() + This is a specialization of paramiko.HostKeys that interfaces with + a storage backend to get the list of host keys. + """ + def __init__(self, storage): + paramiko.HostKeys.__init__(self) - self.namespace = namespace + self.storage = storage - filename = self.get_host_keys_filename() + def load(self, filename): + """Loads all known host keys from the storage backend.""" + self._entries = [] - if os.path.exists(filename): - self.load_host_keys(filename) + lines = self.storage.read_host_keys() - def get_host_keys_filename(self): - """Returns the path to the known host keys file.""" - return os.path.join(self.get_ssh_dir(), 'known_hosts') + for line in lines: + entry = HostKeyEntry.from_line(line) - def get_ssh_dir(self, ssh_dir_name=None): - """Returns the path to the SSH directory on the system. + if entry is not None: + self._entries.append(entry) + + def save(self, filename): + pass - By default, this will attempt to find either a .ssh or ssh directory. - If ``ssh_dir_name`` is specified, the search will be skipped, and we'll - use that name instead. - """ - path = SSHClient._ssh_dir - if not SSHClient._ssh_dir or ssh_dir_name: - path = os.path.expanduser('~') +class SSHClient(paramiko.SSHClient): + """A client for communicating with an SSH server. - if not ssh_dir_name: - ssh_dir_name = '.ssh' + SSHClient allows for communicating with an SSH server and managing + all known host and user keys. - for name in ('.ssh', 'ssh'): - if os.path.exists(os.path.join(path, name)): - ssh_dir_name = name - break + This is a specialization of paramiko.SSHClient, and supports all the + same capabilities. - path = os.path.join(path, ssh_dir_name) + Key access goes through an SSHStorage backend. The storage backend knows + how to look up keys and write them. - if not ssh_dir_name: - SSHClient._ssh_dir = path + The default backend works with the site directory's data/.ssh directory, + and supports namespaced directories for LocalSites. + """ + DEFAULT_STORAGE = 'reviewboard.ssh.storage.FileSSHStorage' + SUPPORTED_KEY_TYPES = (paramiko.RSAKey, paramiko.DSSKey) - if self.namespace: - return os.path.join(path, self.namespace) - else: - return path + def __init__(self, namespace=None, storage=None): + super(SSHClient, self).__init__() + + self.namespace = namespace + self._load_storage() + self._host_keys = SSHHostKeys(self.storage) + + self.load_host_keys('') + + def _load_storage(self): + """Loads the storage backend. + + This will attempt to load the SSH storage backend. If there is an + error in loading the backend, it will be logged, and an + ImproperlyConfigured exception will be raised. + """ + try: + path = getattr(settings, 'RBSSH_STORAGE_BACKEND', + self.DEFAULT_STORAGE) + except ImportError: + # We may not be running in the Django environment. + path = self.DEFAULT_STORAGE + + i = path.rfind('.') + module, class_name = path[:i], path[i + 1:] + + try: + mod = __import__(module, {}, {}, [class_name]) + except ImportError, e: + msg = 'Error importing SSH storage backend %s: "%s"' % (module, e) + logging.critical(msg) + raise ImproperlyConfigured(msg) + + try: + self.storage = getattr(mod, class_name)(namespace=self.namespace) + except Exception, e: + msg = 'Error instantiating SSH storage backend %s: "%s"' % \ + (module, e) + logging.critical(msg) + raise def get_user_key(self): """Returns the keypair of the user running Review Board. @@ -60,33 +104,23 @@ class SSHClient(paramiko.SSHClient): This will be an instance of :py:mod:`paramiko.PKey`, representing a DSS or RSA key, as long as one exists. Otherwise, it may return None. """ - keyfiles = [] + key = None + fp = None - for cls, filename in ((paramiko.RSAKey, 'id_rsa'), - (paramiko.DSSKey, 'id_dsa')): - # Paramiko looks in ~/.ssh and ~/ssh, depending on the platform, - # so check both. - for sshdir in ('.ssh', 'ssh'): - path = os.path.join(self.get_ssh_dir(sshdir), filename) + try: + key = self.storage.read_user_key() + except paramiko.SSHException, e: + logging.error('SSH: Unknown error accessing user key: %s' % e) + except paramiko.PasswordRequiredException, e: + logging.error('SSH: Unable to access password protected ' + 'key file: %s' % e) + except IOError, e: + logging.error('SSH: Error reading user key: %s' % e) - if os.path.isfile(path): - keyfiles.append((cls, path)) + if fp: + fp.close() - for cls, keyfile in keyfiles: - try: - return cls.from_private_key_file(keyfile) - except paramiko.SSHException, e: - logging.error('SSH: Unknown error accessing local key file ' - '%s: %s' - % (keyfile, e)) - except paramiko.PasswordRequiredException, e: - logging.error('SSH: Unable to access password protected ' - 'key file %s: %s' % (keyfile, e)) - except IOError, e: - logging.error('SSH: Error reading local key file %s: %s' - % (keyfile, e)) - - return None + return key def get_public_key(self, key): """Returns the public key portion of an SSH key. @@ -106,78 +140,42 @@ class SSHClient(paramiko.SSHClient): def is_key_authorized(self, key): """Returns whether or not a public key is currently authorized.""" - authorized = False public_key = key.get_base64() try: - filename = os.path.join(self.get_ssh_dir(), 'authorized_keys') - fp = open(filename, 'r') + lines = self.storage.read_authorized_keys() - for line in fp.xreadlines(): + for line in lines: try: authorized_key = line.split()[1] except (ValueError, IndexError): continue if authorized_key == public_key: - authorized = True - break - - fp.close() + return True except IOError: pass - return authorized - - def ensure_ssh_dir(self): - """Ensures the existance of the .ssh directory. + return False - If the directory doesn't exist, it will be created. - The full path to the directory will be returned. - - Callers are expected to handle any exceptions. This may raise - IOError for any problems in creating the directory. - """ - sshdir = self.get_ssh_dir() - - if self.namespace: - # The parent will be the .ssh dir. - parent = os.path.dirname(sshdir) - - if not os.path.exists(parent): - try: - os.mkdir(parent, 0700) - except OSError: - raise MakeSSHDirError(parent) - - if not os.path.exists(sshdir): - try: - os.mkdir(sshdir, 0700) - except OSError: - raise MakeSSHDirError(sshdir) - - return sshdir - - def generate_user_key(self): + def generate_user_key(self, bits=2048): """Generates a new RSA keypair for the user running Review Board. - This will store the new key in :file:`$HOME/.ssh/id_rsa` and return the + This will store the new key in the backend storage and return the resulting key as an instance of :py:mod:`paramiko.RSAKey`. - If a key already exists in the id_rsa file, it's returned instead. + If a key already exists, it's returned instead. Callers are expected to handle any exceptions. This may raise IOError for any problems in writing the key file, or paramiko.SSHException for any other problems. """ - sshdir = self.ensure_ssh_dir() - filename = os.path.join(sshdir, 'id_rsa') + key = self.get_user_key() - if os.path.isfile(filename): - return self.get_user_key() + if not key: + key = paramiko.RSAKey.generate(bits) + self._write_user_key(key) - key = paramiko.RSAKey.generate(2048) - key.write_private_key_file(filename) return key def import_user_key(self, keyfile): @@ -185,8 +183,8 @@ class SSHClient(paramiko.SSHClient): ``keyfile`` is expected to be an ``UploadedFile`` or a paramiko ``KeyFile``. If this is a valid key file, it will be saved in - :file:`$HOME/.ssh/`` and the resulting key as an instance of - :py:mod:`paramiko.RSAKey` will be returned. + the storage backend and the resulting key as an instance of + :py:mod:`paramiko.PKey` will be returned. If a key of this name already exists, it will be overwritten. @@ -197,14 +195,11 @@ class SSHClient(paramiko.SSHClient): This will raise UnsupportedSSHKeyError if the uploaded key is not a supported type. """ - sshdir = self.ensure_ssh_dir() - # Try to find out what key this is. - for cls, filename in ((paramiko.RSAKey, 'id_rsa'), - (paramiko.DSSKey, 'id_dsa')): - try: - key = None + for cls in self.SUPPORTED_KEY_TYPES: + key = None + try: if not isinstance(keyfile, paramiko.PKey): keyfile.seek(0) key = cls.from_private_key(keyfile) @@ -216,67 +211,36 @@ class SSHClient(paramiko.SSHClient): continue if key: - key.write_private_key_file(os.path.join(sshdir, filename)) + self._write_user_key(key) return key raise UnsupportedSSHKeyError() def add_host_key(self, hostname, key): """Adds a host key to the known hosts file.""" - self.ensure_ssh_dir() - filename = self.get_host_keys_filename() - - try: - fp = open(filename, 'a') - fp.write('%s %s %s\n' % (hostname, key.get_name(), - key.get_base64())) - fp.close() - except IOError, e: - raise IOError( - _('Unable to write host keys file %(filename)s: %(error)s') % { - 'filename': filename, - 'error': e, - }) + self.storage.add_host_key(hostname, key) def replace_host_key(self, hostname, old_key, new_key): """Replaces a host key in the known hosts file with another. This is used for replacing host keys that have changed. """ - filename = self.get_host_keys_filename() + self.storage.replace_host_key(hostname, old_key, new_key) - if not os.path.exists(filename): - self.add_host_key(hostname, new_key) - return - - try: - fp = open(filename, 'r') - lines = fp.readlines() - fp.close() - - old_key_base64 = old_key.get_base64() - except IOError, e: - raise IOError( - _('Unable to read host keys file %(filename)s: %(error)s') % { - 'filename': filename, - 'error': e, - }) + def _write_user_key(self, key): + """Convenience function to write a user key and check for errors. + Any errors caused as a result of writing a user key will be logged. + """ try: - fp = open(filename, 'w') - - for line in lines: - parts = line.strip().split(" ") - - if parts[-1] == old_key_base64: - parts[-1] = new_key.get_base64() - - fp.write(' '.join(parts) + '\n') - - fp.close() + self.storage.write_user_key(key) + except UnsupportedSSHKeyError, e: + logging.error('Failed to write unknown key type %s' % type(key)) + raise except IOError, e: - raise IOError( - _('Unable to write host keys file %(filename)s: %(error)s') % { - 'filename': filename, - 'error': e, - }) + logging.error('Failed to write SSH user key: %s' % e) + raise + except Exception, e: + logging.error('Unknown error writing SSH user key: %s' % e, + exc_info=1) + raise diff --git a/reviewboard/ssh/storage.py b/reviewboard/ssh/storage.py new file mode 100644 index 0000000000000000000000000000000000000000..c426e2890f37fb557ab0ed875483b3ca41458520 --- /dev/null +++ b/reviewboard/ssh/storage.py @@ -0,0 +1,249 @@ +import logging +import os + +from django.utils.translation import ugettext_lazy as _ +import paramiko + +from reviewboard.ssh.errors import MakeSSHDirError, UnsupportedSSHKeyError + + +class SSHStorage(object): + def __init__(self, namespace=None): + self.namespace = namespace + + def read_user_key(self): + """Reads the user key. + + This will return an instance of :py:mod:`paramiko.PKey` representing + the user key, if one exists. Otherwise, it will return None. + """ + raise NotImplementedError + + def write_user_key(self, key): + """Writes a user key. + + The user key will be stored, and can be accessed later by + read_user_key. + + This will raise UnsupportedSSHKeyError if ``key`` isn't a + :py:mod:`paramiko.RSAKey` or :py:mod:`paramiko.DSSKey`. + + It may also raise :py:mod:`paramiko.SSHException` for key-related + errors. + """ + raise NotImplementedError + + def read_authorized_keys(self): + """Reads a list of authorized keys. + + The authorized keys are returned as a list of raw key data, which + can then be converted into classes as needed. + """ + raise NotImplementedError + + def read_host_keys(self): + """Reads a list of known host keys. + + This known host keys are returned as a list of raw key data, which + can then be converted into classes as needed. + """ + raise NotImplementedError + + def add_host_key(self, hostname, key): + """Adds a known key for a given host. + + This will store a mapping of the key and hostname so that future + access to the server will know the host is legitimate. + """ + raise NotImplementedError + + def replace_host_key(self, hostname, old_key, new_key): + """Replaces a host key in the known hosts list with another. + + This is used for replacing host keys that have changed. + """ + raise NotImplementedError + + +class FileSSHStorage(SSHStorage): + DEFAULT_KEY_FILES = ( + (paramiko.RSAKey, 'id_rsa'), + (paramiko.DSSKey, 'id_dsa'), + ) + + SSH_DIRS = ('.ssh', 'ssh') + + _ssh_dir = None + + def read_user_key(self): + for cls, filename in self.DEFAULT_KEY_FILES: + # Paramiko looks in ~/.ssh and ~/ssh, depending on the platform, + # so check both. + for sshdir in self.SSH_DIRS: + path = os.path.join(self.get_ssh_dir(sshdir), filename) + + if os.path.isfile(path): + return cls.from_private_key_file(path) + + return None + + def write_user_key(self, key): + key_filename = None + + for cls, filename in self.DEFAULT_KEY_FILES: + if isinstance(key, cls): + key_filename = filename + + if not key_filename: + raise UnsupportedSSHKeyError() + + sshdir = self.ensure_ssh_dir() + filename = os.path.join(sshdir, key_filename) + key.write_private_key_file(filename) + + def read_authorized_keys(self): + filename = os.path.join(self.get_ssh_dir(), 'authorized_keys') + + try: + fp = open(filename, 'r') + lines = fp.readlines() + fp.close() + + return lines + except IOError, e: + logging.warning('Unable to read SSH authorized_keys file %s: %s' + % (filename, e)) + raise + + def read_host_keys(self): + filename = self.get_host_keys_filename() + lines = [] + + if os.path.exists(filename): + try: + fp = open(filename, 'r') + + for line in fp.xreadlines(): + line = line.strip() + + if line and line[0] != '#': + lines.append(line) + + fp.close() + except IOError, e: + logging.error('Unable to read host keys file %s: %s' + % (filename, e)) + + return lines + + def add_host_key(self, hostname, key): + self.ensure_ssh_dir() + filename = self.get_host_keys_filename() + + try: + fp = open(filename, 'a') + fp.write('%s %s %s\n' % (hostname, key.get_name(), + key.get_base64())) + fp.close() + except IOError, e: + raise IOError( + _('Unable to write host keys file %(filename)s: %(error)s') % { + 'filename': filename, + 'error': e, + }) + + def replace_host_key(self, hostname, old_key, new_key): + filename = self.get_host_keys_filename() + + if not os.path.exists(filename): + self.add_host_key(hostname, new_key) + return + + try: + fp = open(filename, 'r') + lines = fp.readlines() + fp.close() + + old_key_base64 = old_key.get_base64() + except IOError, e: + raise IOError( + _('Unable to read host keys file %(filename)s: %(error)s') % { + 'filename': filename, + 'error': e, + }) + + try: + fp = open(filename, 'w') + + for line in lines: + parts = line.strip().split(" ") + + if parts[-1] == old_key_base64: + parts[1] = new_key.get_name() + parts[-1] = new_key.get_base64() + + fp.write(' '.join(parts) + '\n') + + fp.close() + except IOError, e: + raise IOError( + _('Unable to write host keys file %(filename)s: %(error)s') % { + 'filename': filename, + 'error': e, + }) + + def get_host_keys_filename(self): + """Returns the path to the known host keys file.""" + return os.path.join(self.get_ssh_dir(), 'known_hosts') + + def get_ssh_dir(self, ssh_dir_name=None): + """Returns the path to the SSH directory on the system. + + By default, this will attempt to find either a .ssh or ssh directory. + If ``ssh_dir_name`` is specified, the search will be skipped, and we'll + use that name instead. + """ + path = self._ssh_dir + + if not path or ssh_dir_name: + path = os.path.expanduser('~') + + if not ssh_dir_name: + ssh_dir_name = None + + for name in self.SSH_DIRS: + if os.path.exists(os.path.join(path, name)): + ssh_dir_name = name + break + + if not ssh_dir_name: + ssh_dir_name = self.SSH_DIRS[0] + + path = os.path.join(path, ssh_dir_name) + + if not ssh_dir_name: + self.__class__._ssh_dir = path + + if self.namespace: + return os.path.join(path, self.namespace) + else: + return path + + def ensure_ssh_dir(self): + """Ensures the existance of the .ssh directory. + + If the directory doesn't exist, it will be created. + The full path to the directory will be returned. + + Callers are expected to handle any exceptions. This may raise + IOError for any problems in creating the directory. + """ + sshdir = self.get_ssh_dir() + + if not os.path.exists(sshdir): + try: + os.makedirs(sshdir, 0700) + except OSError: + raise MakeSSHDirError(sshdir) + + return sshdir diff --git a/reviewboard/ssh/tests.py b/reviewboard/ssh/tests.py index 0c02e56eaaa9076146831485c436c7f846421466..302f89e2a3e0af4be31b5a3c2b3918a96b37e261 100644 --- a/reviewboard/ssh/tests.py +++ b/reviewboard/ssh/tests.py @@ -6,6 +6,8 @@ from django.test import TestCase as DjangoTestCase import paramiko from reviewboard.ssh.client import SSHClient +from reviewboard.ssh.errors import UnsupportedSSHKeyError +from reviewboard.ssh.storage import FileSSHStorage class SSHTestCase(DjangoTestCase): @@ -13,6 +15,13 @@ class SSHTestCase(DjangoTestCase): self.old_home = os.getenv('HOME') self.tempdir = None os.environ['RBSSH_ALLOW_AGENT'] = '0' + FileSSHStorage._ssh_dir = None + + if not hasattr(SSHTestCase, 'key1'): + SSHTestCase.key1 = paramiko.RSAKey.generate(1024) + SSHTestCase.key2 = paramiko.DSSKey.generate(1024) + SSHTestCase.key1_b64 = SSHTestCase.key1.get_base64() + SSHTestCase.key2_b64 = SSHTestCase.key2.get_base64() def tearDown(self): self._set_home(self.old_home) @@ -24,55 +33,130 @@ class SSHTestCase(DjangoTestCase): os.environ['HOME'] = homedir -class SSHClientTests(SSHTestCase): - """Unit tests for SSHClient.""" +class FileSSHStorageTests(SSHTestCase): + """Unit tests for FileSSHStorage.""" def setUp(self): - super(SSHClientTests, self).setUp() + super(FileSSHStorageTests, self).setUp() self.tempdir = tempfile.mkdtemp(prefix='rb-tests-home-') + self._set_home(self.tempdir) def test_get_ssh_dir_with_dot_ssh(self): - """Testing SSHClient.get_ssh_dir with ~/.ssh""" - self._set_home(self.tempdir) + """Testing FileSSHStorage.get_ssh_dir with ~/.ssh""" sshdir = os.path.join(self.tempdir, '.ssh') - client = SSHClient() - self.assertEqual(client.get_ssh_dir(), sshdir) + storage = FileSSHStorage() + self.assertEqual(storage.get_ssh_dir(), sshdir) def test_get_ssh_dir_with_ssh(self): - """Testing SSHClient.get_ssh_dir with ~/ssh""" - self._set_home(self.tempdir) + """Testing FileSSHStorage.get_ssh_dir with ~/ssh""" sshdir = os.path.join(self.tempdir, 'ssh') os.mkdir(sshdir, 0700) - client = SSHClient() - self.assertEqual(client.get_ssh_dir(), sshdir) + storage = FileSSHStorage() + self.assertEqual(storage.get_ssh_dir(), sshdir) def test_get_ssh_dir_with_dot_ssh_and_localsite(self): - """Testing SSHClient.get_ssh_dir with ~/.ssh and localsite""" - self._set_home(self.tempdir) + """Testing FileSSHStorage.get_ssh_dir with ~/.ssh and localsite""" sshdir = os.path.join(self.tempdir, '.ssh', 'site-1') - client = SSHClient(namespace='site-1') - self.assertEqual(client.get_ssh_dir(), sshdir) + storage = FileSSHStorage(namespace='site-1') + self.assertEqual(storage.get_ssh_dir(), sshdir) def test_get_ssh_dir_with_ssh_and_localsite(self): - """Testing SSHClient.get_ssh_dir with ~/ssh and localsite""" - self._set_home(self.tempdir) + """Testing FileSSHStorage.get_ssh_dir with ~/ssh and localsite""" sshdir = os.path.join(self.tempdir, 'ssh') os.mkdir(sshdir, 0700) sshdir = os.path.join(sshdir, 'site-1') - client = SSHClient(namespace='site-1') - self.assertEqual(client.get_ssh_dir(), sshdir) + storage = FileSSHStorage(namespace='site-1') + self.assertEqual(storage.get_ssh_dir(), sshdir) + + def test_write_user_key_unsupported(self): + """Testing FileSSHStorage.write_user_key with unsupported key type""" + class FakeKey(object): + pass + + storage = FileSSHStorage() + self.assertRaises(UnsupportedSSHKeyError, + lambda: storage.write_user_key(FakeKey())) + + def test_read_host_keys(self): + """Testing FileSSHStorage.read_host_keys""" + storage = FileSSHStorage() + storage.ensure_ssh_dir() + + line1 = 'host1 ssh-rsa %s' % self.key1_b64 + line2 = 'host2 ssh-dss %s' % self.key2_b64 + + filename = storage.get_host_keys_filename() + fp = open(filename, 'w') + fp.write('%s\n' % line1) + fp.write('\n') + fp.write('# foo\n') + fp.write('%s \n' % line2) + fp.close() + + lines = storage.read_host_keys() + self.assertEqual(len(lines), 2) + self.assertEqual(lines[0], line1) + self.assertEqual(lines[1], line2) + + def test_add_host_key(self): + """Testing FileSSHStorage.add_host_key""" + storage = FileSSHStorage() + storage.add_host_key('host1', self.key1) + + filename = storage.get_host_keys_filename() + fp = open(filename, 'r') + lines = fp.readlines() + fp.close() + + self.assertEqual(len(lines), 1) + self.assertEqual(lines[0], 'host1 ssh-rsa %s\n' % self.key1_b64) + + def test_replace_host_key(self): + """Testing FileSSHStorage.replace_host_key""" + storage = FileSSHStorage() + storage.add_host_key('host1', self.key1) + storage.replace_host_key('host1', self.key1, self.key2) + + filename = storage.get_host_keys_filename() + fp = open(filename, 'r') + lines = fp.readlines() + fp.close() + + self.assertEqual(len(lines), 1) + self.assertEqual(lines[0], 'host1 ssh-dss %s\n' % self.key2_b64) + + def test_replace_host_key_no_known_hosts(self): + """Testing FileSSHStorage.replace_host_key with no known hosts file""" + storage = FileSSHStorage() + storage.replace_host_key('host1', self.key1, self.key2) + + filename = storage.get_host_keys_filename() + fp = open(filename, 'r') + lines = fp.readlines() + fp.close() + + self.assertEqual(len(lines), 1) + self.assertEqual(lines[0], 'host1 ssh-dss %s\n' % self.key2_b64) + + +class SSHClientTests(SSHTestCase): + """Unit tests for SSHClient.""" + def setUp(self): + super(SSHClientTests, self).setUp() + + self.tempdir = tempfile.mkdtemp(prefix='rb-tests-home-') def test_generate_user_key(self, namespace=None): """Testing SSHClient.generate_user_key""" self._set_home(self.tempdir) client = SSHClient(namespace=namespace) - key = client.generate_user_key() - key_file = os.path.join(client.get_ssh_dir(), 'id_rsa') + key = client.generate_user_key(bits=1024) + key_file = os.path.join(client.storage.get_ssh_dir(), 'id_rsa') self.assertTrue(os.path.exists(key_file)) self.assertEqual(client.get_user_key(), key) @@ -85,10 +169,9 @@ class SSHClientTests(SSHTestCase): self._set_home(self.tempdir) client = SSHClient(namespace=namespace) - key = paramiko.RSAKey.generate(2048) - client.add_host_key('example.com', key) + client.add_host_key('example.com', self.key1) - known_hosts_file = client.get_host_keys_filename() + known_hosts_file = client.storage.get_host_keys_filename() self.assertTrue(os.path.exists(known_hosts_file)) f = open(known_hosts_file, 'r') @@ -97,7 +180,7 @@ class SSHClientTests(SSHTestCase): self.assertEqual(len(lines), 1) self.assertEqual(lines[0].split(), - ['example.com', key.get_name(), key.get_base64()]) + ['example.com', self.key1.get_name(), self.key1_b64]) def test_add_host_key_with_localsite(self): """Testing SSHClient.add_host_key with localsite""" @@ -108,13 +191,10 @@ class SSHClientTests(SSHTestCase): self._set_home(self.tempdir) client = SSHClient(namespace=namespace) - key = paramiko.RSAKey.generate(2048) - client.add_host_key('example.com', key) + client.add_host_key('example.com', self.key1) + client.replace_host_key('example.com', self.key1, self.key2) - new_key = paramiko.RSAKey.generate(2048) - client.replace_host_key('example.com', key, new_key) - - known_hosts_file = client.get_host_keys_filename() + known_hosts_file = client.storage.get_host_keys_filename() self.assertTrue(os.path.exists(known_hosts_file)) f = open(known_hosts_file, 'r') @@ -123,9 +203,21 @@ class SSHClientTests(SSHTestCase): self.assertEqual(len(lines), 1) self.assertEqual(lines[0].split(), - ['example.com', new_key.get_name(), - new_key.get_base64()]) + ['example.com', self.key2.get_name(), + self.key2_b64]) def test_replace_host_key_with_localsite(self): """Testing SSHClient.replace_host_key with localsite""" self.test_replace_host_key('site-1') + + def test_import_user_key(self, namespace=None): + """Testing SSHClient.import_user_key""" + self._set_home(self.tempdir) + client = SSHClient(namespace=namespace) + + client.import_user_key(self.key1) + self.assertEqual(client.get_user_key(), self.key1) + + def test_import_user_key_with_localsite(self): + """Testing SSHClient.import_user_key with localsite""" + self.test_import_user_key('site-1')