Cleanup this and add refactoring around large constructors (add a parse method). Handle error cases better...

1 files changed, 155 insertions, 122 deletions

diff --git a/cloudinit/ssh_util.py b/cloudinit/ssh_util.py
index 1483f718..93fd55dd 100644
--- a/cloudinit/ssh_util.py
+++ b/cloudinit/ssh_util.py
@@ -20,42 +20,70 @@
 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import os
-import os.path
-import cloudinit.util as util
-
-
-class AuthKeyEntry():
-    # lines are options, keytype, base64-encoded key, comment
-    # man page says the following which I did not understand:
-    #   The options field is optional; its presence is determined by whether
-    #   the line starts with a number or not (the options field never starts
-    #   with a number)
-    options = None
-    keytype = None
-    base64 = None
-    comment = None
-    is_comment = False
-    line_in = ""
+import pwd
+
+from cloudinit import log as logging
+from cloudinit import util
+
+LOG = logging.getLogger(__name__)
+
+
+class AuthKeyEntry(object):
+    """
+    AUTHORIZED_KEYS FILE FORMAT
+     AuthorizedKeysFile specifies the file containing public keys for public
+     key authentication; if none is specified, the default is
+     ~/.ssh/authorized_keys.  Each line of the file contains one key (empty
+     (because of the size of the public key encoding) up to a limit of 8 kilo-
+     bytes, which permits DSA keys up to 8 kilobits and RSA keys up to 16
+     kilobits.  You don't want to type them in; instead, copy the
+     identity.pub, id_dsa.pub, or the id_rsa.pub file and edit it.
+
+     sshd enforces a minimum RSA key modulus size for protocol 1 and protocol
+     2 keys of 768 bits.
+
+     The options (if present) consist of comma-separated option specifica-
+     tions.  No spaces are permitted, except within double quotes.  The fol-
+     lowing option specifications are supported (note that option keywords are
+     case-insensitive):
+    """
 
     def __init__(self, line, def_opt=None):
-        line = line.rstrip("\n\r")
-        self.line_in = line
-        if line.startswith("#") or line.strip() == "":
-            self.is_comment = True
+        self.line = str(line)
+        (self.value, self.components) = self._parse(self.line, def_opt)
+
+    def _form_components(self, toks):
+        components = {}
+        if len(toks) == 1:
+            components['base64'] = toks[0]
+        elif len(toks) == 2:
+            components['base64'] = toks[0]
+            components['comment'] = toks[1]
+        elif len(toks) == 3:
+            components['keytype'] = toks[0]
+            components['base64'] = toks[1]
+            components['comment'] = toks[2]
+        return components
+
+    def get(self, piece):
+        return self.components.get(piece)
+
+    def _parse(self, in_line, def_opt):
+        line = in_line.rstrip("\r\n")
+        if line.startswith("#") or line.strip() == '':
+            return (False, {})
         else:
             ent = line.strip()
             toks = ent.split(None, 3)
-            if len(toks) == 1:
-                self.base64 = toks[0]
-            elif len(toks) == 2:
-                (self.base64, self.comment) = toks
-            elif len(toks) == 3:
-                (self.keytype, self.base64, self.comment) = toks
-            elif len(toks) == 4:
+            tmp_components = {}
+            if def_opt:
+                tmp_components['options'] = def_opt
+            if len(toks) < 4:
+                tmp_components.update(self._form_components(toks))
+            else:
+                # taken from auth_rsa_key_allowed in auth-rsa.c
                 i = 0
-                ent = line.strip()
                 quoted = False
-                # taken from auth_rsa_key_allowed in auth-rsa.c
                 try:
                     while (i < len(ent) and
                            ((quoted) or (ent[i] not in (" ", "\t")))):
@@ -67,124 +95,129 @@ class AuthKeyEntry():
                             quoted = not quoted
                         i = i + 1
                 except IndexError:
-                    self.is_comment = True
-                    return
-
+                    return (False, {})
                 try:
-                    self.options = ent[0:i]
-                    (self.keytype, self.base64, self.comment) = \
-                        ent[i + 1:].split(None, 3)
-                except ValueError:
-                    # we did not understand this line
-                    self.is_comment = True
-
-        if self.options == None and def_opt:
-            self.options = def_opt
-
-        return
-
-    def debug(self):
-        print("line_in=%s\ncomment: %s\noptions=%s\nkeytype=%s\nbase64=%s\n"
-              "comment=%s\n" % (self.line_in, self.is_comment, self.options,
-                                self.keytype, self.base64, self.comment)),
-
-    def __repr__(self):
-        if self.is_comment:
-            return(self.line_in)
+                    options = ent[0:i]
+                    toks = ent[i + 1:].split(None, 3)
+                    if options:
+                        tmp_components['options'] = options
+                    tmp_components.update(self._form_components(toks))
+                except (IndexError, ValueError):
+                    return (False, {})
+            # We got some useful value!
+            return (True, tmp_components)
+
+    def __str__(self):
+        if not self.value:
+            return self.line
         else:
             toks = []
-            for e in (self.options, self.keytype, self.base64, self.comment):
-                if e:
-                    toks.append(e)
-
-            return(' '.join(toks))
+            if 'options' in self.components:
+                toks.append(self.components['options'])
+            if 'keytype' in self.components:
+                toks.append(self.components['keytype'])
+            if 'base64' in self.components:
+                toks.append(self.components['base64'])
+            if 'comment' in self.components:
+                toks.append(self.components['comment'])
+            if not toks:
+                return ''
+            return ' '.join(toks)
 
 
 def update_authorized_keys(fname, keys):
-    # keys is a list of AuthKeyEntries
-    # key_prefix is the prefix (options) to prepend
+    lines = []
     try:
-        fp = open(fname, "r")
-        lines = fp.readlines()  # lines have carriage return
-        fp.close()
-    except IOError:
+        if os.path.isfile(fname):
+            lines = util.load_file(fname).splitlines()
+    except (IOError, OSError):
+        LOG.exception("Error reading lines from %s", fname)
         lines = []
 
-    ka_stats = {}  # keys_added status
-    for k in keys:
-        ka_stats[k] = False
-
-    to_add = []
-    for key in keys:
-        to_add.append(key)
-
+    to_add = list(keys)
     for i in range(0, len(lines)):
         ent = AuthKeyEntry(lines[i])
+        if not ent.value:
+            continue
+        # Replace those with the same base64
         for k in keys:
-            if k.base64 == ent.base64 and not k.is_comment:
+            if not k.value:
+                continue
+            if k.get('base64') == ent.get('base64'):
+                # Replace it with our better one
                 ent = k
-                try:
-                    to_add.remove(k)
-                except ValueError:
-                    pass
+                # Don't add it later
+                to_add.remove(k)
         lines[i] = str(ent)
 
-    # now append any entries we did not match above
+    # Now append any entries we did not match above
     for key in to_add:
         lines.append(str(key))
 
-    if len(lines) == 0:
-        return("")
-    else:
-        return('\n'.join(lines) + "\n")
+    # Ensure it ends with a newline
+    lines.append('')
+    return '\n'.join(lines)
 
 
-def setup_user_keys(keys, user, key_prefix, log=None):
-    import pwd
-    saved_umask = os.umask(077)
-
+def setup_user_keys(keys, user, key_prefix, sshd_config_fn="/etc/ssh/sshd_config"):
     pwent = pwd.getpwnam(user)
 
-    ssh_dir = '%s/.ssh' % pwent.pw_dir
+    ssh_dir = os.path.join(pwent.pw_dir, '.ssh')
     if not os.path.exists(ssh_dir):
-        os.mkdir(ssh_dir)
-        os.chown(ssh_dir, pwent.pw_uid, pwent.pw_gid)
-
-    try:
-        ssh_cfg = parse_ssh_config()
-        akeys = ssh_cfg.get("AuthorizedKeysFile", "%h/.ssh/authorized_keys")
-        akeys = akeys.replace("%h", pwent.pw_dir)
-        akeys = akeys.replace("%u", user)
-        if not akeys.startswith('/'):
-            akeys = os.path.join(pwent.pw_dir, akeys)
-        authorized_keys = akeys
-    except Exception:
-        authorized_keys = '%s/.ssh/authorized_keys' % pwent.pw_dir
-        if log:
-            util.logexc(log)
+        util.ensure_dir(ssh_dir, mode=0700)
+        util.chownbyid(ssh_dir, pwent.pw_uid, pwent.pw_gid)
 
     key_entries = []
     for k in keys:
-        ke = AuthKeyEntry(k, def_opt=key_prefix)
-        key_entries.append(ke)
-
-    content = update_authorized_keys(authorized_keys, key_entries)
-    util.write_file(authorized_keys, content, 0600)
-
-    os.chown(authorized_keys, pwent.pw_uid, pwent.pw_gid)
-    util.restorecon_if_possible(ssh_dir, recursive=True)
-
-    os.umask(saved_umask)
-
-
-def parse_ssh_config(fname="/etc/ssh/sshd_config"):
+        key_entries.append(AuthKeyEntry(k, def_opt=key_prefix))
+
+    with util.SeLinuxGuard(ssh_dir, recursive=True):
+        try:
+            """
+             AuthorizedKeysFile may contain tokens
+             of the form %T which are substituted during connection set-up.
+             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_config_fn)
+            akeys = ssh_cfg.get("authorizedkeysfile", '')
+            akeys = akeys.strip()
+            if not akeys:
+                akeys = "%h/.ssh/authorized_keys"
+            akeys = akeys.replace("%h", pwent.pw_dir)
+            akeys = akeys.replace("%u", user)
+            akeys = akeys.replace("%%", '%')
+            if not akeys.startswith('/'):
+                akeys = os.path.join(pwent.pw_dir, akeys)
+            authorized_keys = akeys
+        except (IOError, OSError):
+            authorized_keys = os.path.join(ssh_dir, 'authorized_keys')
+            LOG.exception(("Failed extracting 'AuthorizedKeysFile' in ssh config"
+                            " from %s, using 'AuthorizedKeysFile' file %s instead."),
+                           sshd_config_fn, authorized_keys)
+
+        content = update_authorized_keys(authorized_keys, key_entries)
+        util.ensure_dir(os.path.dirname(authorized_keys), mode=0700)
+        util.write_file(authorized_keys, content, mode=0600)
+        util.chownbyid(authorized_keys, pwent.pw_uid, pwent.pw_gid)
+
+
+def parse_ssh_config(fname):
+    """
+    The file contains keyword-argu-ment 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 = {}
-    fp = open(fname)
-    for l in fp.readlines():
-        l = l.strip()
-        if not l or l.startswith("#"):
+    if not os.path.isfile(fname):
+        return ret
+    for line in util.load_file(fname).splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
             continue
-        key, val = l.split(None, 1)
-        ret[key] = val
-    fp.close()
-    return(ret)
+        (key, val) = line.split(None, 1)
+        key = key.strip().lower()
+        if key:
+            ret[key] = val
+    return ret