rework package mirror selection

There are several changes here.
 * Datasource now has a 'availability_zone' getter.
 * get_package_mirror_info
   * Datasource convenience 'get_package_mirror_info' that calls
     the configured distro, and passes it the availability-zone
   * distro has a get_package_mirror_info method
   * get_package_mirror_info returns a dict that of name:mirror
     this is to facilitate use of 'security' and 'primary' archive.
   * this supports searching based on templates. Any template
     that references undefined values is skipped. These templates
     can contain 'availability_zone' (LP: #1037727)
   * distro's mirrors can be arch specific (LP: #1028501)
   * 'ec2_region' substitution variable is present only if this
     appears (by availability-zone naming convention) to be in EC2.
 * rename_apt_lists supports the "mirror_info" rather than single mirror
 * generate_sources_list supports mirror_info, and as a result, the
   ubuntu mirrors reference '$security' rather than security (LP: #1006963)
 * remove the DataSourceEc2 specific mirror selection, but instead
   rely on the above filtering, and the fact that 'ec2_region' is only
   defined if the availability_zone looks like a ec2 az.
 * package_mirrors is now required in the system_info config, a dict like:
   package_mirrors:
     - arches: [i386, amd64]
       failsafe:
         primary: http://archive.ubuntu.com/ubuntu
         security: http://security.ubuntu.com/ubuntu
       search:
         primary:
           - http://%(ec2_region)s.ec2.archive.ubuntu.com/ubuntu/
           - http://%(availability_zone)s.clouds.archive.ubuntu.com/ubuntu/
         security: []
     - arches: [armhf, armel, default]
       failsafe:
         primary: http://ports.ubuntu.com/ubuntu
         security: http://ports.ubuntu.com/ubuntu

LP: #1006963, #1028501, #1037727

11 files changed, 303 insertions, 102 deletions

diff --git a/ChangeLog b/ChangeLog
index d6dfc80e..e9e88318 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,4 +1,8 @@
 0.7.0:
+ - allow distro mirror selection to include availability-zone (LP: #1037727)
+ - allow arch specific mirror selection (select ports.ubuntu.com on arm)
+   LP: #1028501
+ - allow specification of security mirrors (LP: #1006963)
  - add the 'None' datasource (LP: #906669), which will allow jobs
    to run even if there is no "real" datasource found.
  - write ssh authorized keys to console, ssh_authkey_fingerprints
diff --git a/cloudinit/cloud.py b/cloudinit/cloud.py
index 22d9167e..620b3c07 100644
--- a/cloudinit/cloud.py
+++ b/cloudinit/cloud.py
@@ -82,9 +82,6 @@ class Cloud(object):
     def get_locale(self):
         return self.datasource.get_locale()
 
-    def get_local_mirror(self):
-        return self.datasource.get_local_mirror()
-
     def get_hostname(self, fqdn=False):
         return self.datasource.get_hostname(fqdn=fqdn)
 
diff --git a/cloudinit/config/cc_apt_update_upgrade.py b/cloudinit/config/cc_apt_update_upgrade.py
index e60e1037..49a93415 100644
--- a/cloudinit/config/cc_apt_update_upgrade.py
+++ b/cloudinit/config/cc_apt_update_upgrade.py
@@ -50,20 +50,25 @@ def handle(name, cfg, cloud, log, _args):
     upgrade = util.get_cfg_option_bool(cfg, 'apt_upgrade', False)
 
     release = get_release()
-    mirror = find_apt_mirror(cloud, cfg)
-    if not mirror:
+    mirrors = find_apt_mirror_info(cloud, cfg)
+    if not mirrors or "primary" not in mirrors:
         log.debug(("Skipping module named %s,"
                    " no package 'mirror' located"), name)
         return
 
-    log.debug("Selected mirror at: %s" % mirror)
+    # backwards compatibility
+    mirror = mirrors["primary"]
+    mirrors["mirror"] = mirror
+
+    log.debug("mirror info: %s" % mirrors)
 
     if not util.get_cfg_option_bool(cfg,
                                     'apt_preserve_sources_list', False):
-        generate_sources_list(release, mirror, cloud, log)
-        old_mir = util.get_cfg_option_str(cfg, 'apt_old_mirror',
-                                          "archive.ubuntu.com/ubuntu")
-        rename_apt_lists(old_mir, mirror)
+        generate_sources_list(release, mirrors, cloud, log)
+        old_mirrors = cfg.get('apt_old_mirrors',
+                              {"primary": "archive.ubuntu.com/ubuntu",
+                               "security": "security.ubuntu.com/ubuntu"})
+        rename_apt_lists(old_mirrors, mirrors)
 
     # Set up any apt proxy
     proxy = cfg.get("apt_proxy", None)
@@ -81,8 +86,10 @@ def handle(name, cfg, cloud, log, _args):
 
     # Process 'apt_sources'
     if 'apt_sources' in cfg:
-        errors = add_sources(cloud, cfg['apt_sources'],
-                             {'MIRROR': mirror, 'RELEASE': release})
+        params = mirrors
+        params['RELEASE'] = release
+        params['MIRROR'] = mirror
+        errors = add_sources(cloud, cfg['apt_sources'], params)
         for e in errors:
             log.warn("Source Error: %s", ':'.join(e))
 
@@ -146,15 +153,18 @@ def mirror2lists_fileprefix(mirror):
     return string
 
 
-def rename_apt_lists(omirror, new_mirror, lists_d="/var/lib/apt/lists"):
-    oprefix = os.path.join(lists_d, mirror2lists_fileprefix(omirror))
-    nprefix = os.path.join(lists_d, mirror2lists_fileprefix(new_mirror))
-    if oprefix == nprefix:
-        return
-    olen = len(oprefix)
-    for filename in glob.glob("%s_*" % oprefix):
-        # TODO(harlowja) use the cloud.paths.join...
-        util.rename(filename, "%s%s" % (nprefix, filename[olen:]))
+def rename_apt_lists(old_mirrors, new_mirrors, lists_d="/var/lib/apt/lists"):
+    for (name, omirror) in old_mirrors.iteritems():
+        nmirror = new_mirrors.get(name)
+        if not nmirror:
+            continue
+        oprefix = os.path.join(lists_d, mirror2lists_fileprefix(omirror))
+        nprefix = os.path.join(lists_d, mirror2lists_fileprefix(nmirror))
+        if oprefix == nprefix:
+            continue
+        olen = len(oprefix)
+        for filename in glob.glob("%s_*" % oprefix):
+            util.rename(filename, "%s%s" % (nprefix, filename[olen:]))
 
 
 def get_release():
@@ -162,14 +172,17 @@ def get_release():
     return stdout.strip()
 
 
-def generate_sources_list(codename, mirror, cloud, log):
+def generate_sources_list(codename, mirrors, cloud, log):
     template_fn = cloud.get_template_filename('sources.list')
-    if template_fn:
-        params = {'mirror': mirror, 'codename': codename}
-        out_fn = cloud.paths.join(False, '/etc/apt/sources.list')
-        templater.render_to_file(template_fn, out_fn, params)
-    else:
+    if not template_fn:
         log.warn("No template found, not rendering /etc/apt/sources.list")
+        return
+
+    params = {'codename': codename}
+    for k in mirrors:
+        params[k] = mirrors[k]
+    out_fn = cloud.paths.join(False, '/etc/apt/sources.list')
+    templater.render_to_file(template_fn, out_fn, params)
 
 
 def add_sources(cloud, srclist, template_params=None):
@@ -231,43 +244,47 @@ def add_sources(cloud, srclist, template_params=None):
     return errorlist
 
 
-def find_apt_mirror(cloud, cfg):
+def find_apt_mirror_info(cloud, cfg):
     """find an apt_mirror given the cloud and cfg provided."""
 
     mirror = None
 
-    cfg_mirror = cfg.get("apt_mirror", None)
-    if cfg_mirror:
-        mirror = cfg["apt_mirror"]
-    elif "apt_mirror_search" in cfg:
-        mirror = util.search_for_mirror(cfg['apt_mirror_search'])
-    else:
-        mirror = cloud.get_local_mirror()
+    # this is less preferred way of specifying mirror preferred would be to
+    # use the distro's search or package_mirror.
+    mirror = cfg.get("apt_mirror", None)
 
-        mydom = ""
+    search = cfg.get("apt_mirror_search", None)
+    if not mirror and search:
+        mirror = util.search_for_mirror(search)
 
+    if (not mirror and
+        util.get_cfg_option_bool(cfg, "apt_mirror_search_dns", False)):
+        mydom = ""
         doms = []
 
-        if not mirror:
-            # if we have a fqdn, then search its domain portion first
-            (_hostname, fqdn) = util.get_hostname_fqdn(cfg, cloud)
-            mydom = ".".join(fqdn.split(".")[1:])
-            if mydom:
-                doms.append(".%s" % mydom)
+        # if we have a fqdn, then search its domain portion first
+        (_hostname, fqdn) = util.get_hostname_fqdn(cfg, cloud)
+        mydom = ".".join(fqdn.split(".")[1:])
+        if mydom:
+            doms.append(".%s" % mydom)
+
+        doms.extend((".localdomain", "",))
 
-        if (not mirror and
-            util.get_cfg_option_bool(cfg, "apt_mirror_search_dns", False)):
-            doms.extend((".localdomain", "",))
+        mirror_list = []
+        distro = cloud.distro.name
+        mirrorfmt = "http://%s-mirror%s/%s" % (distro, "%s", distro)
+        for post in doms:
+            mirror_list.append(mirrorfmt % (post))
 
-            mirror_list = []
-            distro = cloud.distro.name
-            mirrorfmt = "http://%s-mirror%s/%s" % (distro, "%s", distro)
-            for post in doms:
-                mirror_list.append(mirrorfmt % (post))
+        mirror = util.search_for_mirror(mirror_list)
 
-            mirror = util.search_for_mirror(mirror_list)
+    mirror_info = cloud.datasource.get_package_mirror_info()
 
-    if not mirror:
-        mirror = cloud.distro.get_package_mirror()
+    # this is a bit strange.
+    # if mirror is set, then one of the legacy options above set it
+    # but they do not cover security. so we need to get that from
+    # get_package_mirror_info
+    if mirror:
+        mirror_info.update({'primary': mirror})
 
-    return mirror
+    return mirror_info
diff --git a/cloudinit/distros/__init__.py b/cloudinit/distros/__init__.py
index b9609b7a..357209a4 100644
--- a/cloudinit/distros/__init__.py
+++ b/cloudinit/distros/__init__.py
@@ -23,6 +23,8 @@
 from StringIO import StringIO
 
 import abc
+import os
+import re
 
 from cloudinit import importer
 from cloudinit import log as logging
@@ -75,8 +77,26 @@ class Distro(object):
     def update_package_sources(self):
         raise NotImplementedError()
 
-    def get_package_mirror(self):
-        return self.get_option('package_mirror')
+    def get_primary_arch(self):
+        arch = os.uname[4]
+        if arch in ("i386", "i486", "i586", "i686"):
+            return "i386"
+        return arch
+
+    def _get_arch_package_mirror_info(self, arch=None):
+        mirror_info = self.get_option("package_mirrors", None)
+        if arch == None:
+            arch = self.get_primary_arch()
+        return _get_arch_package_mirror_info(mirror_info, arch)
+
+    def get_package_mirror_info(self, arch=None,
+                                availability_zone=None):
+        # this resolves the package_mirrors config option
+        # down to a single dict of {mirror_name: mirror_url}
+        arch_info = self._get_arch_package_mirror_info(arch)
+
+        return _get_package_mirror_info(availability_zone=availability_zone,
+                                        mirror_info=arch_info)
 
     def apply_network(self, settings, bring_up=True):
         # Write it out
@@ -151,6 +171,55 @@ class Distro(object):
             return False
 
 
+def _get_package_mirror_info(mirror_info, availability_zone=None,
+                             mirror_filter=util.search_for_mirror):
+    # given a arch specific 'mirror_info' entry (from package_mirrors)
+    # search through the 'search' entries, and fallback appropriately
+    # return a dict with only {name: mirror} entries.
+
+    ec2_az_re = ("^[a-z][a-z]-(%s)-[1-9][0-9]*[a-z]$" %
+        "north|northeast|east|southeast|south|southwest|west|northwest")
+
+    subst = {}
+    if availability_zone:
+        subst['availability_zone'] = availability_zone
+
+    if availability_zone and re.match(ec2_az_re, availability_zone):
+        subst['ec2_region'] = "%s" % availability_zone[0:-1]
+
+    results = {}
+    for (name, mirror) in mirror_info.get('failsafe', {}).iteritems():
+        results[name] = mirror
+
+    for (name, searchlist) in mirror_info.get('search', {}).iteritems():
+        mirrors = []
+        for tmpl in searchlist:
+            try:
+                mirrors.append(tmpl % subst)
+            except KeyError:
+                pass
+
+        found = mirror_filter(mirrors)
+        if found:
+            results[name] = found
+
+    LOG.debug("filtered distro mirror info: %s" % results)
+
+    return results
+
+
+def _get_arch_package_mirror_info(package_mirrors, arch):
+    # pull out the specific arch from a 'package_mirrors' config option
+    default = None
+    for item in package_mirrors:
+        arches = item.get("arches")
+        if arch in arches:
+            return item
+        if "default" in arches:
+            default = item
+    return default
+
+
 def fetch(name):
     locs = importer.find_module(name,
                                 ['', __name__],
diff --git a/cloudinit/distros/debian.py b/cloudinit/distros/debian.py
index 3247d7ce..da8c1a5b 100644
--- a/cloudinit/distros/debian.py
+++ b/cloudinit/distros/debian.py
@@ -147,3 +147,7 @@ class Distro(distros.Distro):
     def update_package_sources(self):
         self._runner.run("update-sources", self.package_command,
                          ["update"], freq=PER_INSTANCE)
+
+    def get_primary_arch(self):
+        (arch, _err) = util.subp(['dpkg', '--print-architecture'])
+        return str(arch).strip()
diff --git a/cloudinit/sources/DataSourceCloudStack.py b/cloudinit/sources/DataSourceCloudStack.py
index 8056dcfa..f7ffa7cb 100644
--- a/cloudinit/sources/DataSourceCloudStack.py
+++ b/cloudinit/sources/DataSourceCloudStack.py
@@ -131,7 +131,8 @@ class DataSourceCloudStack(sources.DataSource):
     def get_instance_id(self):
         return self.metadata['instance-id']
 
-    def get_availability_zone(self):
+    @property
+    def availability_zone(self):
         return self.metadata['availability-zone']
 
 
diff --git a/cloudinit/sources/DataSourceEc2.py b/cloudinit/sources/DataSourceEc2.py
index d9eb8f17..556dcafb 100644
--- a/cloudinit/sources/DataSourceEc2.py
+++ b/cloudinit/sources/DataSourceEc2.py
@@ -83,40 +83,6 @@ class DataSourceEc2(sources.DataSource):
     def get_availability_zone(self):
         return self.metadata['placement']['availability-zone']
 
-    def get_local_mirror(self):
-        return self.get_mirror_from_availability_zone()
-
-    def get_mirror_from_availability_zone(self, availability_zone=None):
-        # Return type None indicates there is no cloud specific mirror
-        # Availability is like 'us-west-1b' or 'eu-west-1a'
-        if availability_zone is None:
-            availability_zone = self.get_availability_zone()
-
-        if self.is_vpc():
-            return None
-
-        if not availability_zone:
-            return None
-
-        mirror_tpl = self.distro.get_option('package_mirror_ec2_template',
-                                             None)
-
-        if mirror_tpl is None:
-            return None
-
-        # in EC2, the 'region' is 'us-east-1' if 'zone' is 'us-east-1a'
-        tpl_params = {
-            'zone': availability_zone.strip(),
-            'region': availability_zone[:-1]
-        }
-        mirror_url = mirror_tpl % (tpl_params)
-
-        found = util.search_for_mirror([mirror_url])
-        if found is not None:
-            return mirror_url
-
-        return None
-
     def _get_url_settings(self):
         mcfg = self.ds_cfg
         if not mcfg:
@@ -255,6 +221,12 @@ class DataSourceEc2(sources.DataSource):
             return True
         return False
 
+    @property
+    def availability_zone(self):
+        try:
+            return self.metadata['placement']['availability-zone']
+        except KeyError:
+            return None
 
 # Used to match classes to dependencies
 datasources = [
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index ca9f58e5..4719d254 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -117,9 +117,9 @@ class DataSource(object):
     def get_locale(self):
         return 'en_US.UTF-8'
 
-    def get_local_mirror(self):
-        # ??
-        return None
+    @property
+    def availability_zone(self):
+        return self.metadata.get('availability-zone')
 
     def get_instance_id(self):
         if not self.metadata or 'instance-id' not in self.metadata:
@@ -166,6 +166,10 @@ class DataSource(object):
         else:
             return hostname
 
+    def get_package_mirror_info(self):
+        return self.distro.get_package_mirror_info(
+            availability_zone=self.availability_zone)
+
 
 def find_source(sys_cfg, distro, paths, ds_deps, cfg_list, pkg_list):
     ds_list = list_sources(cfg_list, ds_deps, pkg_list)
diff --git a/config/cloud.cfg b/config/cloud.cfg
index 700f3d7a..106ab01a 100644
--- a/config/cloud.cfg
+++ b/config/cloud.cfg
@@ -74,6 +74,18 @@ system_info:
       cloud_dir: /var/lib/cloud/
       templates_dir: /etc/cloud/templates/
       upstart_dir: /etc/init/
-   package_mirror: http://archive.ubuntu.com/ubuntu
-   package_mirror_ec2_template: http://%(region)s.ec2.archive.ubuntu.com/ubuntu/
+   package_mirrors:
+     - arches: [i386, amd64]
+       failsafe:
+         primary: http://archive.ubuntu.com/ubuntu
+         security: http://security.ubuntu.com/ubuntu
+       search:
+         primary:
+           - http://%(ec2_region)s.ec2.archive.ubuntu.com/ubuntu/
+           - http://%(availability_zone)s.clouds.archive.ubuntu.com/ubuntu/
+         security: []
+     - arches: [armhf, armel, default]
+       failsafe:
+         primary: http://ports.ubuntu.com/ubuntu
+         security: http://ports.ubuntu.com/ubuntu
    ssh_svcname: ssh
diff --git a/templates/sources.list.tmpl b/templates/sources.list.tmpl
index f702025f..ce395b3d 100644
--- a/templates/sources.list.tmpl
+++ b/templates/sources.list.tmpl
@@ -52,9 +52,9 @@ deb-src $mirror $codename-updates universe
 # deb http://archive.canonical.com/ubuntu $codename partner
 # deb-src http://archive.canonical.com/ubuntu $codename partner
 
-deb http://security.ubuntu.com/ubuntu $codename-security main
-deb-src http://security.ubuntu.com/ubuntu $codename-security main
-deb http://security.ubuntu.com/ubuntu $codename-security universe
-deb-src http://security.ubuntu.com/ubuntu $codename-security universe
-# deb http://security.ubuntu.com/ubuntu $codename-security multiverse
-# deb-src http://security.ubuntu.com/ubuntu $codename-security multiverse
+deb $security $codename-security main
+deb-src $security $codename-security main
+deb $security $codename-security universe
+deb-src $security $codename-security universe
+# deb $security $codename-security multiverse
+# deb-src $security $codename-security multiverse
diff --git a/tests/unittests/test_distros/test_generic.py b/tests/unittests/test_distros/test_generic.py
new file mode 100644
index 00000000..2df4c2f0
--- /dev/null
+++ b/tests/unittests/test_distros/test_generic.py
@@ -0,0 +1,121 @@
+from mocker import MockerTestCase
+
+from cloudinit import distros
+
+unknown_arch_info = {
+    'arches': ['default'],
+    'failsafe': {'primary': 'http://fs-primary-default',
+                 'security': 'http://fs-security-default'}
+}
+
+package_mirrors = [
+    {'arches': ['i386', 'amd64'],
+     'failsafe': {'primary': 'http://fs-primary-intel',
+                  'security': 'http://fs-security-intel'},
+     'search': {
+         'primary': ['http://%(ec2_region)s.ec2/',
+                     'http://%(availability_zone)s.clouds/'],
+         'security': ['http://security-mirror1-intel',
+                      'http://security-mirror2-intel']}},
+    {'arches': ['armhf', 'armel'],
+     'failsafe': {'primary': 'http://fs-primary-arm',
+                  'security': 'http://fs-security-arm'}},
+    unknown_arch_info
+]
+
+gpmi = distros._get_package_mirror_info  # pylint: disable=W0212
+gapmi = distros._get_arch_package_mirror_info  # pylint: disable=W0212
+
+
+class TestGenericDistro(MockerTestCase):
+
+    def return_first(self, mlist):
+        if not mlist:
+            return None
+        return mlist[0]
+
+    def return_second(self, mlist):
+        if not mlist:
+            return None
+        return mlist[1]
+
+    def return_none(self, _mlist):
+        return None
+
+    def return_last(self, mlist):
+        if not mlist:
+            return None
+        return(mlist[-1])
+
+    def setUp(self):
+        super(TestGenericDistro, self).setUp()
+        # Make a temp directoy for tests to use.
+        self.tmp = self.makeDir()
+
+    def test_arch_package_mirror_info_unknown(self):
+        """for an unknown arch, we should get back that with arch 'default'."""
+        arch_mirrors = gapmi(package_mirrors, arch="unknown")
+        self.assertEqual(unknown_arch_info, arch_mirrors)
+
+    def test_arch_package_mirror_info_known(self):
+        arch_mirrors = gapmi(package_mirrors, arch="amd64")
+        self.assertEqual(package_mirrors[0], arch_mirrors)
+
+    def test_get_package_mirror_info_az_ec2(self):
+        arch_mirrors = gapmi(package_mirrors, arch="amd64")
+
+        results = gpmi(arch_mirrors, availability_zone="us-east-1a",
+                       mirror_filter=self.return_first)
+        self.assertEqual(results,
+                         {'primary': 'http://us-east-1.ec2/',
+                          'security': 'http://security-mirror1-intel'})
+
+        results = gpmi(arch_mirrors, availability_zone="us-east-1a",
+                       mirror_filter=self.return_second)
+        self.assertEqual(results,
+                         {'primary': 'http://us-east-1a.clouds/',
+                          'security': 'http://security-mirror2-intel'})
+
+        results = gpmi(arch_mirrors, availability_zone="us-east-1a",
+                       mirror_filter=self.return_none)
+        self.assertEqual(results, package_mirrors[0]['failsafe'])
+
+    def test_get_package_mirror_info_az_non_ec2(self):
+        arch_mirrors = gapmi(package_mirrors, arch="amd64")
+
+        results = gpmi(arch_mirrors, availability_zone="nova.cloudvendor",
+                       mirror_filter=self.return_first)
+        self.assertEqual(results,
+                         {'primary': 'http://nova.cloudvendor.clouds/',
+                          'security': 'http://security-mirror1-intel'})
+
+        results = gpmi(arch_mirrors, availability_zone="nova.cloudvendor",
+                       mirror_filter=self.return_last)
+        self.assertEqual(results,
+                         {'primary': 'http://nova.cloudvendor.clouds/',
+                          'security': 'http://security-mirror2-intel'})
+
+    def test_get_package_mirror_info_none(self):
+        arch_mirrors = gapmi(package_mirrors, arch="amd64")
+
+        # because both search entries here replacement based on
+        # availability-zone, the filter will be called with an empty list and
+        # failsafe should be taken.
+        results = gpmi(arch_mirrors, availability_zone=None,
+                       mirror_filter=self.return_first)
+        self.assertEqual(results,
+                         {'primary': 'http://fs-primary-intel',
+                          'security': 'http://security-mirror1-intel'})
+
+        results = gpmi(arch_mirrors, availability_zone=None,
+                       mirror_filter=self.return_last)
+        self.assertEqual(results,
+                         {'primary': 'http://fs-primary-intel',
+                          'security': 'http://security-mirror2-intel'})
+
+
+#def _get_package_mirror_info(mirror_info, availability_zone=None,
+#                             mirror_filter=util.search_for_mirror):
+
+
+# vi: ts=4 expandtab