From 8dffe8af24604e1949cc991c1db181a69695d945 Mon Sep 17 00:00:00 2001
From: Bastian Blank <waldi@debian.org>
Date: Tue, 16 Aug 2022 15:45:11 +0200
Subject: [PATCH] config: Support APT automated mirror selection
Forwarded: https://github.com/canonical/cloud-init/pull/1670

---
 cloudinit/config/cc_apt_configure.py          | 22 +++++-
 .../schemas/schema-cloud-config-v1.json       |  5 ++
 .../test_apt_configure_mirrorlists_v3.py      | 68 +++++++++++++++++++
 3 files changed, 94 insertions(+), 1 deletion(-)
 create mode 100644 tests/unittests/config/test_apt_configure_mirrorlists_v3.py

Index: cloud-init/cloudinit/config/cc_apt_configure.py
===================================================================
--- cloud-init.orig/cloudinit/config/cc_apt_configure.py
+++ cloud-init/cloudinit/config/cc_apt_configure.py
@@ -220,7 +220,10 @@ def apply_apt(cfg, cloud, target):
     mirrors = find_apt_mirror_info(cfg, cloud, arch=arch)
     LOG.debug("Apt Mirror info: %s", mirrors)
 
-    if util.is_false(cfg.get("preserve_sources_list", False)):
+    if util.is_true(cfg.get("generate_mirrorlists", False)):
+        generate_mirrorlists(cfg, mirrors, cloud)
+
+    elif util.is_false(cfg.get("preserve_sources_list", False)):
         add_mirror_keys(cfg, target)
         generate_sources_list(cfg, release, mirrors, cloud)
         rename_apt_lists(mirrors, target, arch)
@@ -485,6 +488,23 @@ def generate_sources_list(cfg, release,
     util.write_file(aptsrc, disabled, mode=0o644)
 
 
+def generate_mirrorlists(cfg, mirrors, cloud):
+    """generate_mirrorlists
+    create one file for every mirror for apt-transport-mirror(1)"""
+    aptmir = pathlib.Path("/etc/apt/mirrors")
+    util.ensure_dir(str(aptmir))
+    util.write_file(
+        str(aptmir / f"{cloud.distro.name}.list"),
+        f"{mirrors['PRIMARY']}\n",
+        mode=0o644,
+    )
+    util.write_file(
+        str(aptmir / f"{cloud.distro.name}-security.list"),
+        f"{mirrors['SECURITY']}\n",
+        mode=0o644,
+    )
+
+
 def add_apt_key_raw(key, file_name, hardened=False, target=None):
     """
     actual adding of a key as defined in key argument
Index: cloud-init/cloudinit/config/schemas/schema-cloud-config-v1.json
===================================================================
--- cloud-init.orig/cloudinit/config/schemas/schema-cloud-config-v1.json
+++ cloud-init/cloudinit/config/schemas/schema-cloud-config-v1.json
@@ -749,6 +749,11 @@
               "default": false,
               "description": "By default, cloud-init will generate a new sources list in ``/etc/apt/sources.list.d`` based on any changes specified in cloud config. To disable this behavior and preserve the sources list from the pristine image, set ``preserve_sources_list`` to ``true``.\n\nThe ``preserve_sources_list`` option overrides all other config keys that would alter ``sources.list`` or ``sources.list.d``, **except** for additional sources to be added to ``sources.list.d``."
             },
+            "generate_mirrorlists": {
+              "type": "boolean",
+              "default": false,
+              "description": "Write lists for APT automated mirror selection (``apt-transport-mirror(1)``). It will write separate lists for both the PRIMARY and SECURITY mirror into ``/etc/apt/mirrors/${DIST}.list`` and ``/etc/apt/mirrors/${DIST}-security.list``.  Those can then be used in the APT source.list as ``mirror+file:///etc/apt/mirrors/${DIST}.list``. No ``/etc/apt/sources.list`` will be writte in this case."
+            },
             "disable_suites": {
               "type": "array",
               "items": {
Index: cloud-init/tests/unittests/config/test_apt_configure_mirrorlists_v3.py
===================================================================
--- /dev/null
+++ cloud-init/tests/unittests/config/test_apt_configure_mirrorlists_v3.py
@@ -0,0 +1,68 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+""" test_apt_custom_mirrorlists
+Test creation of mirrorlists
+"""
+import logging
+import shutil
+import tempfile
+from contextlib import ExitStack
+from unittest import mock
+from unittest.mock import call
+
+from cloudinit import subp, util
+from cloudinit.config import cc_apt_configure
+from tests.unittests import helpers as t_help
+from tests.unittests.util import get_cloud
+
+LOG = logging.getLogger(__name__)
+
+
+class TestAptSourceConfigMirrorlists(t_help.FilesystemMockingTestCase):
+    """TestAptSourceConfigMirrorlists - Class to test mirrorlists rendering"""
+
+    def setUp(self):
+        super().setUp()
+        self.subp = subp.subp
+        self.new_root = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.new_root)
+
+        rpatcher = mock.patch("cloudinit.util.lsb_release")
+        get_rel = rpatcher.start()
+        get_rel.return_value = {"codename": "fakerel"}
+        self.addCleanup(rpatcher.stop)
+        apatcher = mock.patch("cloudinit.util.get_dpkg_architecture")
+        get_arch = apatcher.start()
+        get_arch.return_value = "amd64"
+        self.addCleanup(apatcher.stop)
+
+    def test_apt_v3_mirrors_list(self):
+        """test_apt_v3_mirrors_list"""
+        cfg = {"apt": {"generate_mirrorlists": True}}
+
+        mycloud = get_cloud("ubuntu")
+
+        with ExitStack() as stack:
+            mock_writefile = stack.enter_context(
+                mock.patch.object(util, "write_file")
+            )
+            stack.enter_context(mock.patch.object(util, "ensure_dir"))
+            cc_apt_configure.handle("test", cfg, mycloud, LOG, None)
+
+        mock_writefile.assert_has_calls(
+            [
+                call(
+                    "/etc/apt/mirrors/ubuntu.list",
+                    "http://archive.ubuntu.com/ubuntu/\n",
+                    mode=0o644,
+                ),
+                call(
+                    "/etc/apt/mirrors/ubuntu-security.list",
+                    "http://security.ubuntu.com/ubuntu/\n",
+                    mode=0o644,
+                ),
+            ]
+        )
+
+
+# vi: ts=4 expandtab