Simplify sshd config loading to make it use one set of functions/objects/classes

2 files changed, 62 insertions, 32 deletions

diff --git a/cloudinit/config/cc_set_passwords.py b/cloudinit/config/cc_set_passwords.py
index eb68ddfe..5b72224b 100644
--- a/cloudinit/config/cc_set_passwords.py
+++ b/cloudinit/config/cc_set_passwords.py
@@ -20,6 +20,7 @@
 
 import sys
 
+from cloudinit import ssh_util
 from cloudinit import util
 
 from string import letters, digits  # pylint: disable=W0402
@@ -101,39 +102,31 @@ def handle(_name, cfg, cloud, log, args):
             pw_auth = 'no'
 
     if change_pwauth:
-        new_lines = []
         replaced_auth = False
-        replacement = "PasswordAuthentication %s" % (pw_auth)
-
-        # See http://linux.die.net/man/5/sshd_config
-        conf_fn = cloud.paths.join(True, '/etc/ssh/sshd_config')
-        # Todo: use the common ssh_util function for this parsing...
-        old_lines = util.load_file(conf_fn).splitlines()
-        for i, line in enumerate(old_lines):
-            if not line.strip() or line.startswith("#"):
-                new_lines.append(line)
-                continue
-            splitup = line.split(None, 1)
-            if len(splitup) <= 1:
-                new_lines.append(line)
-                continue
-            (cmd, args) = splitup
+
+        # See: man sshd_config
+        conf_fn = cloud.paths.join(True, ssh_util.DEF_SSHD_CFG)
+        old_lines = ssh_util.parse_ssh_config(conf_fn)
+        new_lines = []
+        i = 0
+        for (i, line) in enumerate(old_lines):
             # Keywords are case-insensitive and arguments are case-sensitive
-            cmd = cmd.lower().strip()
-            if cmd == 'passwordauthentication':
-                log.debug("Replacing auth line %s with %s", i + 1, replacement)
+            if line.key == 'passwordauthentication':
+                log.debug("Replacing auth line %s with %s", i + 1, pw_auth)
                 replaced_auth = True
-                new_lines.append(replacement)
-            else:
-                new_lines.append(line)
+                line.value = pw_auth
+            new_lines.append(line)
 
         if not replaced_auth:
-            log.debug("Adding new auth line %s", replacement)
+            log.debug("Adding new auth line %s", i + 1)
             replaced_auth = True
-            new_lines.append(replacement)
+            new_lines.append(ssh_util.SshdConfigLine('',
+                                                     'PasswordAuthentication',
+                                                     pw_auth))
 
-        util.write_file(cloud.paths.join(False, '/etc/ssh/sshd_config'),
-                        "\n".join(new_lines))
+        lines = [str(e) for e in new_lines]
+        ssh_rw_fn = cloud.paths.join(False, ssh_util.DEF_SSHD_CFG)
+        util.write_file(ssh_rw_fn, "\n".join(lines))
 
         try:
             cmd = ['service']
diff --git a/cloudinit/ssh_util.py b/cloudinit/ssh_util.py
index 45dd5535..fc8b9b3d 100644
--- a/cloudinit/ssh_util.py
+++ b/cloudinit/ssh_util.py
@@ -29,6 +29,8 @@ from cloudinit import log as logging
 from cloudinit import util
 
 LOG = logging.getLogger(__name__)
+
+# See: man sshd_config
 DEF_SSHD_CFG = "/etc/ssh/sshd_config"
 
 
@@ -233,7 +235,7 @@ def setup_user_keys(keys, user, key_prefix, paths):
             # The following tokens are defined: %% is replaced by a literal
             # '%', %h is replaced by the home directory of the user being
             # authenticated and %u is replaced by the username of that user.
-            ssh_cfg = parse_ssh_config(sshd_conf_fn)
+            ssh_cfg = parse_ssh_config_map(sshd_conf_fn)
             akeys = ssh_cfg.get("authorizedkeysfile", '')
             akeys = akeys.strip()
             if not akeys:
@@ -258,19 +260,54 @@ def setup_user_keys(keys, user, key_prefix, paths):
         util.chownbyid(authorized_keys, pwent.pw_uid, pwent.pw_gid)
 
 
+class SshdConfigLine(object):
+    def __init__(self, line, k=None, v=None):
+        self.line = line
+        self._key = k
+        self.value = v
+
+    @property
+    def key(self):
+        if self._key is None:
+            return None
+        # Keywords are case-insensitive
+        return self._key.lower()
+
+    def __str__(self):
+        if self._key is None:
+            return str(self.line)
+        else:
+            v = str(self._key)
+            if self.value:
+                v += " " + str(self.value)
+            return v
+
+
 def parse_ssh_config(fname):
+    # See: man sshd_config
     # The file contains keyword-argument pairs, one per line.
     # Lines starting with '#' and empty lines are interpreted as comments.
     # Note: key-words are case-insensitive and arguments are case-sensitive
-    ret = {}
+    lines = []
     if not os.path.isfile(fname):
-        return ret
+        return lines
     for line in util.load_file(fname).splitlines():
         line = line.strip()
         if not line or line.startswith("#"):
+            lines.append(SshdConfigLine(line))
             continue
         (key, val) = line.split(None, 1)
-        key = key.strip().lower()
-        if key:
-            ret[key] = val
+        lines.append(SshdConfigLine(line, key, val))
+    return lines
+
+
+def parse_ssh_config_map(fname):
+    lines = parse_ssh_config(fname)
+    if not lines:
+        return {}
+    ret = {}
+    for line in lines:
+        if not line.key:
+            continue
+        ret[line.key] = line.value
     return ret