From 0496cb8db99addf879af32c76d06952ee7374416 Mon Sep 17 00:00:00 2001
From: vitstransky <vit.stransky@deutsche-boerse.com>
Date: Fri, 27 Sep 2024 13:46:47 +0200
Subject: [PATCH] API Key for authentication (#149)

* API Key implementation

* Update artifactory_cleanup/cli.py

---------

Co-authored-by: allburov <allburov@gmail.com>
---
 artifactory_cleanup/cli.py         | 12 ++++++++++--
 artifactory_cleanup/loaders.py     | 17 ++++++++++-------
 tests/data/all-built-in-rules.yaml |  1 +
 tests/test_loaders.py              |  4 +++-
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/artifactory_cleanup/cli.py b/artifactory_cleanup/cli.py
index 999744b..98edb54 100644
--- a/artifactory_cleanup/cli.py
+++ b/artifactory_cleanup/cli.py
@@ -166,9 +166,17 @@ def main(self):
             print(str(err), file=sys.stderr)
             sys.exit(1)
 
-        server, user, password = loader.get_connection()
+        server, user, password, apikey = loader.get_connection()
         session = BaseUrlSession(server)
-        session.auth = HTTPBasicAuth(user, password)
+        if apikey:
+            print("Using API Key")
+            headers = {
+                "X-JFrog-Art-Api": apikey
+            }
+            session.headers = headers
+        else:
+            print("Using user and password")
+            session.auth = HTTPBasicAuth(user, password)
 
         self._destroy_or_verbose()
         print(f"Using {self._worker_count} workers")
diff --git a/artifactory_cleanup/loaders.py b/artifactory_cleanup/loaders.py
index d96b47d..63aff57 100644
--- a/artifactory_cleanup/loaders.py
+++ b/artifactory_cleanup/loaders.py
@@ -94,10 +94,11 @@ def get_root_schema(self, rules):
         config_schema = cfgv.Map(
             "Config",
             None,
-            cfgv.NoAdditionalKeys(["server", "user", "password", "policies"]),
+            cfgv.NoAdditionalKeys(["server", "user", "password", "policies", "apikey"]),
             cfgv.Required("server", cfgv.check_string),
-            cfgv.Required("user", cfgv.check_string),
-            cfgv.Required("password", cfgv.check_string),
+            # User and password required, if apikey missing
+            cfgv.Conditional("user", cfgv.check_string, "apikey", cfgv.MISSING, False),
+            cfgv.Conditional("password", cfgv.check_string, "apikey", cfgv.MISSING, False),
             cfgv.RequiredRecurse("policies", cfgv.Array(policy_schema)),
         )
 
@@ -184,16 +185,18 @@ def load(filename):
             filename, schema, yaml.safe_load, InvalidConfigError
         )
 
-    def get_connection(self) -> Tuple[str, str, str]:
+    def get_connection(self) -> Tuple[str, str, str, str]:
         config = self.load(self.filepath)
         server = config["artifactory-cleanup"]["server"]
-        user = config["artifactory-cleanup"]["user"]
-        password = config["artifactory-cleanup"]["password"]
+        user = config.get("artifactory-cleanup", {}).get("user", "")
+        password = config.get("artifactory-cleanup", {}).get("password", "")
+        apikey = config.get("artifactory-cleanup", {}).get("apikey", "")
 
         server = os.path.expandvars(server)
         user = os.path.expandvars(user)
         password = os.path.expandvars(password)
-        return server, user, password
+        apikey = os.path.expandvars(apikey)
+        return server, user, password, apikey
 
 
 class PythonLoader:
diff --git a/tests/data/all-built-in-rules.yaml b/tests/data/all-built-in-rules.yaml
index b733946..5a7151e 100644
--- a/tests/data/all-built-in-rules.yaml
+++ b/tests/data/all-built-in-rules.yaml
@@ -2,6 +2,7 @@ artifactory-cleanup:
   server: https://repo.example.com/artifactory
   user: $ARTIFACTORY_USERNAME
   password: $ARTIFACTORY_PASSWORD
+  apikey: $ARTIFACTORY_APIKEY
 
   policies:
     - name: repo-without-name
diff --git a/tests/test_loaders.py b/tests/test_loaders.py
index d6d822b..6b29ece 100644
--- a/tests/test_loaders.py
+++ b/tests/test_loaders.py
@@ -10,10 +10,12 @@ def test_all_rules(self, shared_datadir):
     def test_load_env_variables(self, shared_datadir, monkeypatch):
         monkeypatch.setenv("ARTIFACTORY_USERNAME", "UserName")
         monkeypatch.setenv("ARTIFACTORY_PASSWORD", "P@ssw0rd")
+        monkeypatch.setenv("ARTIFACTORY_APIKEY", "Ap1Key")
 
         loader = YamlConfigLoader(shared_datadir / "all-built-in-rules.yaml")
-        server, user, password = loader.get_connection()
+        server, user, password, apikey = loader.get_connection()
 
         assert server == "https://repo.example.com/artifactory"
         assert user == "UserName"
         assert password == "P@ssw0rd"
+        assert apikey == "Ap1Key"