From 8d2d00ae4d16ed3888a2d6be54b23093b5247db6 Mon Sep 17 00:00:00 2001
From: Dusty Mabe <dusty@dustymabe.com>
Date: Jun 05 2020 15:47:40 +0000
Subject: coreos: Detect what key to use based on build id


Our Fedora CoreOS version numbers are predictable in that they
encode the Fedora major version into the first 2 digits of the
version; i.e, `31.20200505.3.0`. Let's auto detect what key to
use based on that build id.

Fixes: https://github.com/coreos/fedora-coreos-tracker/issues/296

Signed-off-by: Dusty Mabe <dusty@dustymabe.com>

---

diff --git a/robosignatory.toml b/robosignatory.toml
index 06efe40..51452ce 100644
--- a/robosignatory.toml
+++ b/robosignatory.toml
@@ -125,6 +125,8 @@ handlers = ["console"]
 
     [consumer_config.coreos]
     bucket = "robosig-dev-fcos-builds"
+    # Only set key if you want to override key detection that detects
+    # which key to use based on FCOS version numbers.
     key = "coreos"
 
         [consumer_config.coreos.aws]
diff --git a/robosignatory/coreos.py b/robosignatory/coreos.py
index 11a367e..6e903b0 100644
--- a/robosignatory/coreos.py
+++ b/robosignatory/coreos.py
@@ -39,7 +39,17 @@ class CoreOSSigner(object):
     def get_key(self, msg):
         # Evaluation of the key is here and not in __init__ because we may want
         # a stream or version-dependant key in the future.
-        return self.config["coreos"]["key"]
+
+        # If there is a key hardcoded in the config then use that.
+        # This is useful in the staging environment where we only have
+        # one key.
+        config_key = self.config["coreos"].get("key")
+        if config_key:
+            return config_key
+        # Detect what key to use by using the first digits of the
+        # build_id (in the form of `"build_id": "32.20200527.20.0"`).
+        major = int(msg.body["build_id"].split('.')[0])
+        return 'fedora-' + str(major)
 
     def consume(self, msg):
         # Message structure:
diff --git a/tests/test_coreos.py b/tests/test_coreos.py
index 84b9cc3..edac1f3 100644
--- a/tests/test_coreos.py
+++ b/tests/test_coreos.py
@@ -172,3 +172,24 @@ class TestCoreOS(unittest.TestCase):
 
         self.consumer.bucket.download_file.assert_called()
         self.consumer.bucket.upload_file.assert_not_called()
+
+    def test_key_parse_config(self):
+        # Verify that when a key is provided via the config it is used
+        consumer = CoreOSSigner(TEST_CONFIG)
+        msg = Message(topic=ARTIFACTS_MESSAGE.topic, body=ARTIFACTS_MESSAGE.body)
+        self.assertEqual(consumer.get_key(msg), "testing")
+
+    def test_key_parse_autodetect(self):
+        # Verify that when no key is provided via the config the key
+        # is autodetected.
+        # 
+        # Grab the config and remove the hardcoded key to enable auto detection
+        config = copy.deepcopy(TEST_CONFIG)
+        del config["coreos"]["key"]
+        consumer = CoreOSSigner(config)
+        # Grab the message body and set the build_id to a version number
+        new_body = copy.deepcopy(ARTIFACTS_MESSAGE.body)
+        new_body["build_id"] = "32.20200601.2.1"
+        # Process the Message and verify the key matches what we expect
+        msg = Message(topic=ARTIFACTS_MESSAGE.topic, body=new_body)
+        self.assertEqual(consumer.get_key(msg), "fedora-32")