Get Azure endpoint server from DHCP client

It is more efficient and cross-distribution safe to use the hooks function
from dhclient to obtain the Azure endpoint server (DHCP option 245).

This is done by providing shell scritps that are called by the hooks
infrastructure of both dhclient and NetworkManager.  The hooks then
invoke 'cloud-init dhclient-hook' that maintains json data
with the dhclient options in
/run/cloud-init/dhclient.hooks/<interface>.json .

The azure helper then pulls the value from
/run/cloud-init/dhclient.hooks/<interface>.json file(s). If that file does
not exist or the value is not present, it will then fall back to the
original method of scraping the dhcp client lease file.

5 files changed, 199 insertions, 35 deletions

diff --git a/cloudinit/atomic_helper.py b/cloudinit/atomic_helper.py
new file mode 100644
index 00000000..15319f71
--- /dev/null
+++ b/cloudinit/atomic_helper.py
@@ -0,0 +1,25 @@
+#!/usr/bin/python
+# vi: ts=4 expandtab
+
+import json
+import os
+import tempfile
+
+
+def atomic_write_file(path, content, mode='w'):
+    tf = None
+    try:
+        tf = tempfile.NamedTemporaryFile(dir=os.path.dirname(path),
+                                         delete=False, mode=mode)
+        tf.write(content)
+        tf.close()
+        os.rename(tf.name, path)
+    except Exception as e:
+        if tf is not None:
+            os.unlink(tf.name)
+        raise e
+
+
+def atomic_write_json(path, data):
+    return atomic_write_file(path, json.dumps(data, indent=1,
+                                              sort_keys=True) + "\n")
diff --git a/cloudinit/cmd/main.py b/cloudinit/cmd/main.py
index 63621c1d..ba22b168 100644
--- a/cloudinit/cmd/main.py
+++ b/cloudinit/cmd/main.py
@@ -25,7 +25,6 @@ import argparse
 import json
 import os
 import sys
-import tempfile
 import time
 import traceback
 
@@ -47,6 +46,10 @@ from cloudinit.reporting import events
 from cloudinit.settings import (PER_INSTANCE, PER_ALWAYS, PER_ONCE,
                                 CLOUD_CONFIG)
 
+from cloudinit.atomic_helper import atomic_write_json
+
+from cloudinit.dhclient_hook import LogDhclient
+
 
 # Pretty little cheetah formatted welcome message template
 WELCOME_MSG_TPL = ("Cloud-init v. ${version} running '${action}' at "
@@ -452,22 +455,10 @@ def main_single(name, args):
         return 0
 
 
-def atomic_write_file(path, content, mode='w'):
-    tf = None
-    try:
-        tf = tempfile.NamedTemporaryFile(dir=os.path.dirname(path),
-                                         delete=False, mode=mode)
-        tf.write(content)
-        tf.close()
-        os.rename(tf.name, path)
-    except Exception as e:
-        if tf is not None:
-            os.unlink(tf.name)
-        raise e
-
-
-def atomic_write_json(path, data):
-    return atomic_write_file(path, json.dumps(data, indent=1) + "\n")
+def dhclient_hook(name, args):
+    record = LogDhclient(args)
+    record.check_hooks_dir()
+    record.record()
 
 
 def status_wrapper(name, args, data_d=None, link_d=None):
@@ -627,7 +618,6 @@ def main(sysv_args=None):
     # This subcommand allows you to run a single module
     parser_single = subparsers.add_parser('single',
                                           help=('run a single module '))
-    parser_single.set_defaults(action=('single', main_single))
     parser_single.add_argument("--name", '-n', action="store",
                                help="module name to run",
                                required=True)
@@ -644,6 +634,16 @@ def main(sysv_args=None):
                                      ' pass to this module'))
     parser_single.set_defaults(action=('single', main_single))
 
+    parser_dhclient = subparsers.add_parser('dhclient-hook',
+                                            help=('run the dhclient hook'
+                                                  'to record network info'))
+    parser_dhclient.add_argument("net_action",
+                                 help=('action taken on the interface'))
+    parser_dhclient.add_argument("net_interface",
+                                 help=('the network interface being acted'
+                                       ' upon'))
+    parser_dhclient.set_defaults(action=('dhclient_hook', dhclient_hook))
+
     args = parser.parse_args(args=sysv_args)
 
     try:
@@ -677,9 +677,18 @@ def main(sysv_args=None):
                         "running single module %s" % args.name)
         report_on = args.report
 
+    elif name == 'dhclient_hook':
+        rname, rdesc = ("dhclient-hook",
+                        "running dhclient-hook module")
+
     args.reporter = events.ReportEventStack(
         rname, rdesc, reporting_enabled=report_on)
+
     with args.reporter:
         return util.log_time(
             logfunc=LOG.debug, msg="cloud-init mode '%s'" % name,
             get_uptime=True, func=functor, args=(name, args))
+
+
+if __name__ == '__main__':
+    main(sys.argv)
diff --git a/cloudinit/dhclient_hook.py b/cloudinit/dhclient_hook.py
new file mode 100644
index 00000000..9dcbe39c
--- /dev/null
+++ b/cloudinit/dhclient_hook.py
@@ -0,0 +1,50 @@
+#!/usr/bin/python
+# vi: ts=4 expandtab
+
+import os
+
+from cloudinit.atomic_helper import atomic_write_json
+from cloudinit import log as logging
+from cloudinit import stages
+
+LOG = logging.getLogger(__name__)
+
+
+class LogDhclient(object):
+
+    def __init__(self, cli_args):
+        self.hooks_dir = self._get_hooks_dir()
+        self.net_interface = cli_args.net_interface
+        self.net_action = cli_args.net_action
+        self.hook_file = os.path.join(self.hooks_dir,
+                                      self.net_interface + ".json")
+
+    @staticmethod
+    def _get_hooks_dir():
+        i = stages.Init()
+        return os.path.join(i.paths.get_runpath(), 'dhclient.hooks')
+
+    def check_hooks_dir(self):
+        if not os.path.exists(self.hooks_dir):
+            os.makedirs(self.hooks_dir)
+        else:
+            # If the action is down and the json file exists, we need to
+            # delete the file
+            if self.net_action is 'down' and os.path.exists(self.hook_file):
+                os.remove(self.hook_file)
+
+    @staticmethod
+    def get_vals(info):
+        new_info = {}
+        for k, v in info.items():
+            if k.startswith("DHCP4_") or k.startswith("new_"):
+                key = (k.replace('DHCP4_', '').replace('new_', '')).lower()
+                new_info[key] = v
+        return new_info
+
+    def record(self):
+        envs = os.environ
+        if self.hook_file is None:
+            return
+        atomic_write_json(self.hook_file, self.get_vals(envs))
+        LOG.debug("Wrote dhclient options in %s", self.hook_file)
diff --git a/cloudinit/sources/DataSourceAzure.py b/cloudinit/sources/DataSourceAzure.py
index 8c7e8673..a251fe01 100644
--- a/cloudinit/sources/DataSourceAzure.py
+++ b/cloudinit/sources/DataSourceAzure.py
@@ -20,18 +20,17 @@ import base64
 import contextlib
 import crypt
 import fnmatch
+from functools import partial
 import os
 import os.path
 import time
-import xml.etree.ElementTree as ET
-
 from xml.dom import minidom
-
-from cloudinit.sources.helpers.azure import get_metadata_from_fabric
+import xml.etree.ElementTree as ET
 
 from cloudinit import log as logging
 from cloudinit.settings import PER_ALWAYS
 from cloudinit import sources
+from cloudinit.sources.helpers.azure import get_metadata_from_fabric
 from cloudinit import util
 
 LOG = logging.getLogger(__name__)
@@ -107,6 +106,8 @@ def temporary_hostname(temp_hostname, cfg, hostname_command='hostname'):
 
 
 class DataSourceAzureNet(sources.DataSource):
+    FALLBACK_LEASE = '/var/lib/dhcp/dhclient.eth0.leases'
+
     def __init__(self, sys_cfg, distro, paths):
         sources.DataSource.__init__(self, sys_cfg, distro, paths)
         self.seed_dir = os.path.join(paths.seed_dir, 'azure')
@@ -115,6 +116,8 @@ class DataSourceAzureNet(sources.DataSource):
         self.ds_cfg = util.mergemanydict([
             util.get_cfg_by_path(sys_cfg, DS_CFG_PATH, {}),
             BUILTIN_DS_CONFIG])
+        self.dhclient_lease_file = self.paths.cfgs.get('dhclient_lease',
+                                                       self.FALLBACK_LEASE)
 
     def __str__(self):
         root = sources.DataSource.__str__(self)
@@ -226,7 +229,9 @@ class DataSourceAzureNet(sources.DataSource):
         write_files(ddir, files, dirmode=0o700)
 
         if self.ds_cfg['agent_command'] == '__builtin__':
-            metadata_func = get_metadata_from_fabric
+            metadata_func = partial(get_metadata_from_fabric,
+                                    fallback_lease_file=self.
+                                    dhclient_lease_file)
         else:
             metadata_func = self.get_metadata_from_agent
         try:
diff --git a/cloudinit/sources/helpers/azure.py b/cloudinit/sources/helpers/azure.py
index 63ccf10e..6e43440f 100644
--- a/cloudinit/sources/helpers/azure.py
+++ b/cloudinit/sources/helpers/azure.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import re
@@ -6,6 +7,7 @@ import struct
 import tempfile
 import time
 
+from cloudinit import stages
 from contextlib import contextmanager
 from xml.etree import ElementTree
 
@@ -187,19 +189,32 @@ class WALinuxAgentShim(object):
         '  </Container>',
         '</Health>'])
 
-    def __init__(self):
+    def __init__(self, fallback_lease_file=None):
         LOG.debug('WALinuxAgentShim instantiated...')
-        self.endpoint = self.find_endpoint()
+        self.dhcpoptions = None
+        self._endpoint = None
         self.openssl_manager = None
         self.values = {}
+        self.lease_file = fallback_lease_file
 
     def clean_up(self):
         if self.openssl_manager is not None:
             self.openssl_manager.clean_up()
 
     @staticmethod
-    def get_ip_from_lease_value(lease_value):
-        unescaped_value = lease_value.replace('\\', '')
+    def _get_hooks_dir():
+        _paths = stages.Init()
+        return os.path.join(_paths.paths.get_runpath(), "dhclient.hooks")
+
+    @property
+    def endpoint(self):
+        if self._endpoint is None:
+            self._endpoint = self.find_endpoint(self.lease_file)
+        return self._endpoint
+
+    @staticmethod
+    def get_ip_from_lease_value(fallback_lease_value):
+        unescaped_value = fallback_lease_value.replace('\\', '')
         if len(unescaped_value) > 4:
             hex_string = ''
             for hex_pair in unescaped_value.split(':'):
@@ -213,15 +228,75 @@ class WALinuxAgentShim(object):
         return socket.inet_ntoa(packed_bytes)
 
     @staticmethod
-    def find_endpoint():
-        LOG.debug('Finding Azure endpoint...')
-        content = util.load_file('/var/lib/dhcp/dhclient.eth0.leases')
-        value = None
+    def _get_value_from_leases_file(fallback_lease_file):
+        leases = []
+        content = util.load_file(fallback_lease_file)
+        LOG.debug("content is {}".format(content))
         for line in content.splitlines():
             if 'unknown-245' in line:
-                value = line.strip(' ').split(' ', 2)[-1].strip(';\n"')
+                # Example line from Ubuntu
+                # option unknown-245 a8:3f:81:10;
+                leases.append(line.strip(' ').split(' ', 2)[-1].strip(';\n"'))
+        # Return the "most recent" one in the list
+        if len(leases) < 1:
+            return None
+        else:
+            return leases[-1]
+
+    @staticmethod
+    def _load_dhclient_json():
+        dhcp_options = {}
+        hooks_dir = WALinuxAgentShim._get_hooks_dir()
+        if not os.path.exists(hooks_dir):
+            LOG.debug("%s not found.", hooks_dir)
+            return None
+        hook_files = [os.path.join(hooks_dir, x)
+                      for x in os.listdir(hooks_dir)]
+        for hook_file in hook_files:
+            try:
+                name = os.path.basename(hook_file).replace('.json', '')
+                dhcp_options[name] = json.loads(util.load_file((hook_file)))
+            except ValueError:
+                raise ValueError("%s is not valid JSON data", hook_file)
+        return dhcp_options
+
+    @staticmethod
+    def _get_value_from_dhcpoptions(dhcp_options):
+        if dhcp_options is None:
+            return None
+        # the MS endpoint server is given to us as DHPC option 245
+        _value = None
+        for interface in dhcp_options:
+            _value = dhcp_options[interface].get('unknown_245', None)
+            if _value is not None:
+                LOG.debug("Endpoint server found in dhclient options")
+                break
+        return _value
+
+    @staticmethod
+    def find_endpoint(fallback_lease_file=None):
+        LOG.debug('Finding Azure endpoint...')
+        value = None
+        # Option-245 stored in /run/cloud-init/dhclient.hooks/<ifc>.json
+        # a dhclient exit hook that calls cloud-init-dhclient-hook
+        dhcp_options = WALinuxAgentShim._load_dhclient_json()
+        value = WALinuxAgentShim._get_value_from_dhcpoptions(dhcp_options)
         if value is None:
-            raise ValueError('No endpoint found in DHCP config.')
+            # Fallback and check the leases file if unsuccessful
+            LOG.debug("Unable to find endpoint in dhclient logs. "
+                      " Falling back to check lease files")
+            if fallback_lease_file is None:
+                LOG.warn("No fallback lease file was specified.")
+                value = None
+            else:
+                LOG.debug("Looking for endpoint in lease file %s",
+                          fallback_lease_file)
+                value = WALinuxAgentShim._get_value_from_leases_file(
+                    fallback_lease_file)
+
+        if value is None:
+            raise ValueError('No endpoint found.')
+
         endpoint_ip_address = WALinuxAgentShim.get_ip_from_lease_value(value)
         LOG.debug('Azure endpoint found at %s', endpoint_ip_address)
         return endpoint_ip_address
@@ -271,8 +346,8 @@ class WALinuxAgentShim(object):
         LOG.info('Reported ready to Azure fabric.')
 
 
-def get_metadata_from_fabric():
-    shim = WALinuxAgentShim()
+def get_metadata_from_fabric(fallback_lease_file=None):
+    shim = WALinuxAgentShim(fallback_lease_file=fallback_lease_file)
     try:
         return shim.register_with_azure_and_fetch_data()
     finally: