There were a couple of ssh-related utilitiy functions scattered in io.py. We are moving them to ssh.py to keep all ssh-related code together.
Signed-off-by: Helga Velroyen <[email protected]> --- lib/backend.py | 2 +- lib/ssh.py | 94 ++++++++++++++++++++++++++++++++++++- lib/tools/prepare_node_join.py | 2 +- lib/utils/io.py | 89 ----------------------------------- test/py/ganeti.ssh_unittest.py | 68 +++++++++++++++++++++++++++ test/py/ganeti.utils.io_unittest.py | 68 --------------------------- 6 files changed, 162 insertions(+), 161 deletions(-) diff --git a/lib/backend.py b/lib/backend.py index f28d004..c72e74b 100644 --- a/lib/backend.py +++ b/lib/backend.py @@ -550,7 +550,7 @@ def LeaveCluster(modify_ssh_setup): try: priv_key, pub_key, auth_keys = ssh.GetUserFiles(constants.SSH_LOGIN_USER) - utils.RemoveAuthorizedKey(auth_keys, utils.ReadFile(pub_key)) + ssh.RemoveAuthorizedKey(auth_keys, utils.ReadFile(pub_key)) utils.RemoveFile(priv_key) utils.RemoveFile(pub_key) diff --git a/lib/ssh.py b/lib/ssh.py index 677c2e4..d54684b 100644 --- a/lib/ssh.py +++ b/lib/ssh.py @@ -24,8 +24,9 @@ """ -import os import logging +import os +import tempfile from ganeti import utils from ganeti import errors @@ -106,6 +107,95 @@ def GetAllUserFiles(user, mkdir=False, dircheck=True, _homedir_fn=None): for (kind, (privkey, pubkey, _)) in result)) +def _SplitSshKey(key): + """Splits a line for SSH's C{authorized_keys} file. + + If the line has no options (e.g. no C{command="..."}), only the significant + parts, the key type and its hash, are used. Otherwise the whole line is used + (split at whitespace). + + @type key: string + @param key: Key line + @rtype: tuple + + """ + parts = key.split() + + if parts and parts[0] in constants.SSHAK_ALL: + # If the key has no options in front of it, we only want the significant + # fields + return (False, parts[:2]) + else: + # Can't properly split the line, so use everything + return (True, parts) + + +def AddAuthorizedKey(file_obj, key): + """Adds an SSH public key to an authorized_keys file. + + @type file_obj: str or file handle + @param file_obj: path to authorized_keys file + @type key: str + @param key: string containing key + + """ + key_fields = _SplitSshKey(key) + + if isinstance(file_obj, basestring): + f = open(file_obj, "a+") + else: + f = file_obj + + try: + nl = True + for line in f: + # Ignore whitespace changes + if _SplitSshKey(line) == key_fields: + break + nl = line.endswith("\n") + else: + if not nl: + f.write("\n") + f.write(key.rstrip("\r\n")) + f.write("\n") + f.flush() + finally: + f.close() + + +def RemoveAuthorizedKey(file_name, key): + """Removes an SSH public key from an authorized_keys file. + + @type file_name: str + @param file_name: path to authorized_keys file + @type key: str + @param key: string containing key + + """ + key_fields = _SplitSshKey(key) + + fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name)) + try: + out = os.fdopen(fd, "w") + try: + f = open(file_name, "r") + try: + for line in f: + # Ignore whitespace changes while comparing lines + if _SplitSshKey(line) != key_fields: + out.write(line) + + out.flush() + os.rename(tmpname, file_name) + finally: + f.close() + finally: + out.close() + except: + utils.RemoveFile(tmpname) + raise + + def InitSSHSetup(error_fn=errors.OpPrereqError): """Setup the SSH configuration for the node. @@ -127,7 +217,7 @@ def InitSSHSetup(error_fn=errors.OpPrereqError): raise error_fn("Could not generate ssh keypair, error %s" % result.output) - utils.AddAuthorizedKey(auth_keys, utils.ReadFile(pub_key)) + AddAuthorizedKey(auth_keys, utils.ReadFile(pub_key)) class SshRunner: diff --git a/lib/tools/prepare_node_join.py b/lib/tools/prepare_node_join.py index c14f410..789c872 100644 --- a/lib/tools/prepare_node_join.py +++ b/lib/tools/prepare_node_join.py @@ -226,7 +226,7 @@ def UpdateSshRoot(data, dry_run, _homedir_fn=None): logging.info("This is a dry run, not modifying %s", auth_keys_file) else: for (_, _, public_key) in keys: - utils.AddAuthorizedKey(auth_keys_file, public_key) + ssh.AddAuthorizedKey(auth_keys_file, public_key) def LoadData(raw): diff --git a/lib/utils/io.py b/lib/utils/io.py index 015cd03..7870a31 100644 --- a/lib/utils/io.py +++ b/lib/utils/io.py @@ -852,95 +852,6 @@ def ReadLockedPidFile(path): return None -def _SplitSshKey(key): - """Splits a line for SSH's C{authorized_keys} file. - - If the line has no options (e.g. no C{command="..."}), only the significant - parts, the key type and its hash, are used. Otherwise the whole line is used - (split at whitespace). - - @type key: string - @param key: Key line - @rtype: tuple - - """ - parts = key.split() - - if parts and parts[0] in constants.SSHAK_ALL: - # If the key has no options in front of it, we only want the significant - # fields - return (False, parts[:2]) - else: - # Can't properly split the line, so use everything - return (True, parts) - - -def AddAuthorizedKey(file_obj, key): - """Adds an SSH public key to an authorized_keys file. - - @type file_obj: str or file handle - @param file_obj: path to authorized_keys file - @type key: str - @param key: string containing key - - """ - key_fields = _SplitSshKey(key) - - if isinstance(file_obj, basestring): - f = open(file_obj, "a+") - else: - f = file_obj - - try: - nl = True - for line in f: - # Ignore whitespace changes - if _SplitSshKey(line) == key_fields: - break - nl = line.endswith("\n") - else: - if not nl: - f.write("\n") - f.write(key.rstrip("\r\n")) - f.write("\n") - f.flush() - finally: - f.close() - - -def RemoveAuthorizedKey(file_name, key): - """Removes an SSH public key from an authorized_keys file. - - @type file_name: str - @param file_name: path to authorized_keys file - @type key: str - @param key: string containing key - - """ - key_fields = _SplitSshKey(key) - - fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name)) - try: - out = os.fdopen(fd, "w") - try: - f = open(file_name, "r") - try: - for line in f: - # Ignore whitespace changes while comparing lines - if _SplitSshKey(line) != key_fields: - out.write(line) - - out.flush() - os.rename(tmpname, file_name) - finally: - f.close() - finally: - out.close() - except: - RemoveFile(tmpname) - raise - - def DaemonPidFileName(name): """Compute a ganeti pid file absolute path diff --git a/test/py/ganeti.ssh_unittest.py b/test/py/ganeti.ssh_unittest.py index eedc852..621d2c0 100755 --- a/test/py/ganeti.ssh_unittest.py +++ b/test/py/ganeti.ssh_unittest.py @@ -148,5 +148,73 @@ class TestGetUserFiles(unittest.TestCase): self.assertEqual(os.listdir(self.tmpdir), []) +class TestSshKeys(testutils.GanetiTestCase): + """Test case for the AddAuthorizedKey function""" + + KEY_A = "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a" + KEY_B = ('command="/usr/bin/fooserver -t --verbose",from="198.51.100.4" ' + "ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b") + + def setUp(self): + testutils.GanetiTestCase.setUp(self) + self.tmpname = self._CreateTempFile() + handle = open(self.tmpname, "w") + try: + handle.write("%s\n" % TestSshKeys.KEY_A) + handle.write("%s\n" % TestSshKeys.KEY_B) + finally: + handle.close() + + def testAddingNewKey(self): + ssh.AddAuthorizedKey(self.tmpname, + "ssh-dss AAAAB3NzaC1kc3MAAACB root@test") + + self.assertFileContent(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" + 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' + " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n" + "ssh-dss AAAAB3NzaC1kc3MAAACB root@test\n") + + def testAddingAlmostButNotCompletelyTheSameKey(self): + ssh.AddAuthorizedKey(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@test") + + # Only significant fields are compared, therefore the key won't be + # updated/added + self.assertFileContent(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" + 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' + " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n") + + def testAddingExistingKeyWithSomeMoreSpaces(self): + ssh.AddAuthorizedKey(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a") + ssh.AddAuthorizedKey(self.tmpname, + "ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22") + + self.assertFileContent(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" + 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' + " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n" + "ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22\n") + + def testRemovingExistingKeyWithSomeMoreSpaces(self): + ssh.RemoveAuthorizedKey(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a") + + self.assertFileContent(self.tmpname, + 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' + " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n") + + def testRemovingNonExistingKey(self): + ssh.RemoveAuthorizedKey(self.tmpname, + "ssh-dss AAAAB3Nsdfj230xxjxJjsjwjsjdjU root@test") + + self.assertFileContent(self.tmpname, + "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" + 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' + " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n") + + if __name__ == "__main__": testutils.GanetiTestProgram() diff --git a/test/py/ganeti.utils.io_unittest.py b/test/py/ganeti.utils.io_unittest.py index 12cc4df..a9232f6 100755 --- a/test/py/ganeti.utils.io_unittest.py +++ b/test/py/ganeti.utils.io_unittest.py @@ -841,74 +841,6 @@ class TestPidFileFunctions(unittest.TestCase): shutil.rmtree(self.dir) -class TestSshKeys(testutils.GanetiTestCase): - """Test case for the AddAuthorizedKey function""" - - KEY_A = "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a" - KEY_B = ('command="/usr/bin/fooserver -t --verbose",from="198.51.100.4" ' - "ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b") - - def setUp(self): - testutils.GanetiTestCase.setUp(self) - self.tmpname = self._CreateTempFile() - handle = open(self.tmpname, "w") - try: - handle.write("%s\n" % TestSshKeys.KEY_A) - handle.write("%s\n" % TestSshKeys.KEY_B) - finally: - handle.close() - - def testAddingNewKey(self): - utils.AddAuthorizedKey(self.tmpname, - "ssh-dss AAAAB3NzaC1kc3MAAACB root@test") - - self.assertFileContent(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" - 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' - " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n" - "ssh-dss AAAAB3NzaC1kc3MAAACB root@test\n") - - def testAddingAlmostButNotCompletelyTheSameKey(self): - utils.AddAuthorizedKey(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@test") - - # Only significant fields are compared, therefore the key won't be - # updated/added - self.assertFileContent(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" - 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' - " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n") - - def testAddingExistingKeyWithSomeMoreSpaces(self): - utils.AddAuthorizedKey(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a") - utils.AddAuthorizedKey(self.tmpname, - "ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22") - - self.assertFileContent(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" - 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' - " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n" - "ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22\n") - - def testRemovingExistingKeyWithSomeMoreSpaces(self): - utils.RemoveAuthorizedKey(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a") - - self.assertFileContent(self.tmpname, - 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' - " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n") - - def testRemovingNonExistingKey(self): - utils.RemoveAuthorizedKey(self.tmpname, - "ssh-dss AAAAB3Nsdfj230xxjxJjsjwjsjdjU root@test") - - self.assertFileContent(self.tmpname, - "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n" - 'command="/usr/bin/fooserver -t --verbose",from="198.51.100.4"' - " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n") - - class TestNewUUID(unittest.TestCase): """Test case for NewUUID""" -- 2.1.0.rc2.206.gedb03e5
