Add in another helper that can understand the

'etc/hosts' format and add in a unit test to
make sure that format can be correctly handled
and added onto in a nice manner + update the 
distro code to use this new code instead of the
previous function that did the same thing.

2 files changed, 118 insertions, 47 deletions

diff --git a/cloudinit/distros/__init__.py b/cloudinit/distros/__init__.py
index 3e9d934d..bf4317f1 100644
--- a/cloudinit/distros/__init__.py
+++ b/cloudinit/distros/__init__.py
@@ -34,6 +34,8 @@ from cloudinit import log as logging
 from cloudinit import ssh_util
 from cloudinit import util
 
+from cloudinit.distros import helpers
+
 # TODO(harlowja): Make this via config??
 IFACE_ACTIONS = {
     'up': ['ifup', '--all'],
@@ -152,42 +154,43 @@ class Distro(object):
         return "127.0.0.1"
 
     def update_etc_hosts(self, hostname, fqdn):
-        # Format defined at
-        # http://unixhelp.ed.ac.uk/CGI/man-cgi?hosts
-        header = "# Added by cloud-init"
-        real_header = "%s on %s" % (header, util.time_rfc2822())
+        header = ''
+        if os.path.exists('/etc/hosts'):
+            eh = helpers.HostsConf(util.load_file("/etc/hosts"))
+        else:
+            eh = helpers.HostsConf('')
+            header = "# Added by cloud-init"
+            header = "%s on %s" % (header, util.time_rfc2822())
         local_ip = self._get_localhost_ip()
-        hosts_line = "%s\t%s %s" % (local_ip, fqdn, hostname)
-        new_etchosts = StringIO()
-        need_write = False
-        need_change = True
-        hosts_ro_fn = self._paths.join(True, "/etc/hosts")
-        for line in util.load_file(hosts_ro_fn).splitlines():
-            if line.strip().startswith(header):
-                continue
-            if not line.strip() or line.strip().startswith("#"):
-                new_etchosts.write("%s\n" % (line))
-                continue
-            split_line = [s.strip() for s in line.split()]
-            if len(split_line) < 2:
-                new_etchosts.write("%s\n" % (line))
-                continue
-            (ip, hosts) = split_line[0], split_line[1:]
-            if ip == local_ip:
-                if sorted([hostname, fqdn]) == sorted(hosts):
-                    need_change = False
-                if need_change:
-                    line = "%s\n%s" % (real_header, hosts_line)
+        prev_info = eh.get_entry(local_ip)
+        need_change = False
+        if not prev_info:
+            eh.add_entry(local_ip, fqdn, hostname)
+            need_change = True
+        else:
+            need_change = True
+            for entry in prev_info:
+                if sorted(entry) == sorted([fqdn, hostname]):
+                    # Exists already, leave it be
                     need_change = False
-                    need_write = True
-            new_etchosts.write("%s\n" % (line))
+                    break
+            if need_change:
+                # Doesn't exist, change the first
+                # entry to be this entry
+                new_entries = list(prev_info)
+                new_entries[0] = [fqdn, hostname]
+                eh.del_entries(local_ip)
+                for entry in new_entries:
+                    if len(entry) == 1:
+                        eh.add_entry(local_ip, entry[0])
+                    elif len(entry) >= 2:
+                        eh.add_entry(local_ip, *entry)
         if need_change:
-            new_etchosts.write("%s\n%s\n" % (real_header, hosts_line))
-            need_write = True
-        if need_write:
-            contents = new_etchosts.getvalue()
-            util.write_file(self._paths.join(False, "/etc/hosts"),
-                            contents, mode=0644)
+            contents = StringIO()
+            if header:
+                contents.write("%s\n" % (header))
+            contents.write("%s\n" % (eh))
+            util.write_file("/etc/hosts", contents.getvalue(), mode=0644)
 
     def _interface_action(self, action):
         if action not in IFACE_ACTIONS:
diff --git a/cloudinit/distros/helpers.py b/cloudinit/distros/helpers.py
index 251d5051..916935eb 100644
--- a/cloudinit/distros/helpers.py
+++ b/cloudinit/distros/helpers.py
@@ -23,6 +23,87 @@ from StringIO import StringIO
 from cloudinit import util
 
 
+def _chop_comment(text, comment_chars):
+    comment_locations = [text.find(c) for c in comment_chars]
+    comment_locations = [c for c in comment_locations if c != -1]
+    if not comment_locations:
+        return (text, '')
+    min_comment = min(comment_locations)
+    before_comment = text[0:min_comment]
+    comment = text[min_comment:]
+    return (before_comment, comment)
+
+
+# See: man hosts
+# or http://unixhelp.ed.ac.uk/CGI/man-cgi?hosts
+class HostsConf(object):
+    def __init__(self, text):
+        self._text = text
+        self._contents = None
+
+    def parse(self):
+        if self._contents is None:
+            self._contents = self._parse(self._text)
+
+    def get_entry(self, ip):
+        self.parse()
+        options = []
+        for (line_type, components) in self._contents:
+            if line_type == 'option':
+                (pieces, _tail) = components
+                if len(pieces) and pieces[0] == ip:
+                    options.append(pieces[1:])
+        return options
+
+    def del_entries(self, ip):
+        self.parse()
+        n_entries = []
+        for (line_type, components) in self._contents:
+            if line_type != 'option':
+                n_entries.append((line_type, components))
+                continue
+            else:
+                (pieces, _tail) = components
+                if len(pieces) and pieces[0] == ip:
+                    pass
+                elif len(pieces):
+                    n_entries.append((line_type, list(components)))
+        self._contents = n_entries
+
+    def add_entry(self, ip, canonical_hostname, *aliases):
+        self.parse()
+        self._contents.append(('option',
+                              ([ip, canonical_hostname] + list(aliases), '')))
+
+    def _parse(self, contents):
+        entries = []
+        for line in contents.splitlines():
+            if not len(line.strip()):
+                entries.append(('blank', [line]))
+                continue
+            (head, tail) = _chop_comment(line.strip(), '#')
+            if not len(head):
+                entries.append(('all_comment', [line]))
+                continue
+            entries.append(('option', [head.split(None), tail]))
+        return entries
+
+    def __str__(self):
+        self.parse()
+        contents = StringIO()
+        for (line_type, components) in self._contents:
+            if line_type == 'blank':
+                contents.write("%s\n")
+            elif line_type == 'all_comment':
+                contents.write("%s\n" % (components[0]))
+            elif line_type == 'option':
+                (pieces, tail) = components
+                pieces = [str(p) for p in pieces]
+                pieces = "\t".join(pieces)
+                contents.write("%s%s\n" % (pieces, tail))
+        return contents.getvalue()
+
+
 # See: man resolv.conf
 class ResolvConf(object):
     def __init__(self, text):
@@ -151,20 +232,7 @@ class ResolvConf(object):
             if not sline:
                 entries.append(('blank', [line]))
                 continue
-            comment_s_loc = sline.find(";")
-            comment_h_loc = sline.find("#")
-            comment_loc = -1
-            if comment_s_loc != -1 and comment_h_loc != -1:
-                comment_loc = min(comment_h_loc, comment_s_loc)
-            elif comment_s_loc != -1:
-                comment_loc = comment_s_loc
-            elif comment_h_loc != -1:
-                comment_loc = comment_h_loc
-            head = line
-            tail = None
-            if comment_loc != -1:
-                head = line[:comment_loc]
-                tail = line[comment_loc:]
+            (head, tail) = _chop_comment(line, ';#')
             if not len(head.strip()):
                 entries.append(('all_comment', [line]))
                 continue