From 9d6de42f4889258510bd3f95fe547c24eee7b97f Mon Sep 17 00:00:00 2001
From: Satoe Imaishi <simaishi@redhat.com>
Date: Mon, 29 Nov 2021 10:03:49 -0500
Subject: [PATCH] Allow rekey with an existing key

(cherry picked from commit 0c6440b46756f02a669d87e461faa4abc5bab8e6)
---
 .../commands/regenerate_secret_key.py         | 17 +++++++++++--
 .../commands/test_secret_key_regeneration.py  | 24 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/awx/main/management/commands/regenerate_secret_key.py b/awx/main/management/commands/regenerate_secret_key.py
index 09d110d0ee..9c86ef4e6e 100644
--- a/awx/main/management/commands/regenerate_secret_key.py
+++ b/awx/main/management/commands/regenerate_secret_key.py
@@ -16,13 +16,26 @@ from awx.main.utils.encryption import encrypt_field, decrypt_field, encrypt_valu
 
 class Command(BaseCommand):
     """
-    Regenerate a new SECRET_KEY value and re-encrypt every secret in the database.
+    Re-encrypt every secret in the database, using regenerated new SECRET_KEY or user provided key.
     """
 
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--use-custom-key',
+            dest='use_custom_key',
+            action='store_true',
+            default=False,
+            help='Use existing key provided as TOWER_SECRET_KEY environment variable',
+        )
+
     @transaction.atomic
     def handle(self, **options):
         self.old_key = settings.SECRET_KEY
-        self.new_key = base64.encodebytes(os.urandom(33)).decode().rstrip()
+        custom_key = os.environ.get("TOWER_SECRET_KEY")
+        if options.get("use_custom_key") and custom_key:
+            self.new_key = custom_key
+        else:
+            self.new_key = base64.encodebytes(os.urandom(33)).decode().rstrip()
         self._notification_templates()
         self._credentials()
         self._unified_jobs()
diff --git a/awx/main/tests/functional/commands/test_secret_key_regeneration.py b/awx/main/tests/functional/commands/test_secret_key_regeneration.py
index c90894c663..6bd0f0a576 100644
--- a/awx/main/tests/functional/commands/test_secret_key_regeneration.py
+++ b/awx/main/tests/functional/commands/test_secret_key_regeneration.py
@@ -3,6 +3,8 @@ import json
 from cryptography.fernet import InvalidToken
 from django.test.utils import override_settings
 from django.conf import settings
+from django.core.management import call_command
+import os
 import pytest
 
 from awx.main import models
@@ -158,3 +160,25 @@ class TestKeyRegeneration:
         # verify that the new SECRET_KEY *does* work
         with override_settings(SECRET_KEY=new_key):
             assert models.OAuth2Application.objects.get(pk=oauth_application.pk).client_secret == secret
+
+    def test_use_custom_key_with_tower_secret_key_env_var(self):
+        custom_key = 'MXSq9uqcwezBOChl/UfmbW1k4op+bC+FQtwPqgJ1u9XV'
+        os.environ['TOWER_SECRET_KEY'] = custom_key
+        new_key = call_command('regenerate_secret_key', '--use-custom-key')
+        assert custom_key == new_key
+
+    def test_use_custom_key_with_empty_tower_secret_key_env_var(self):
+        os.environ['TOWER_SECRET_KEY'] = ''
+        new_key = call_command('regenerate_secret_key', '--use-custom-key')
+        assert settings.SECRET_KEY != new_key
+
+    def test_use_custom_key_with_no_tower_secret_key_env_var(self):
+        os.environ.pop('TOWER_SECRET_KEY', None)
+        new_key = call_command('regenerate_secret_key', '--use-custom-key')
+        assert settings.SECRET_KEY != new_key
+
+    def test_with_tower_secret_key_env_var(self):
+        custom_key = 'MXSq9uqcwezBOChl/UfmbW1k4op+bC+FQtwPqgJ1u9XV'
+        os.environ['TOWER_SECRET_KEY'] = custom_key
+        new_key = call_command('regenerate_secret_key')
+        assert custom_key != new_key