Use JWTs with sync-cli (#434)

* Use JWTs with sync-cli

* Remove unused constant

* Adding a warning message if both were provided

* Updated message to explain what's actually going to happen

Author: michael-richey

datadog_sync/commands/shared/options.py

@@ -42,6 +42,13 @@ def handle_parse_result(self, ctx: Context, opts: Dict[Any, Any], args: List[Any
         help="Datadog source organization APP key.",
         cls=CustomOptionClass,
     ),
+    option(
+        "--source-jwt",
+        envvar=constants.DD_SOURCE_JWT,
+        required=False,
+        help="Datadog source organization JWT (takes precedence over API key).",
+        cls=CustomOptionClass,
+    ),
     option(
         "--source-api-url",
         envvar=constants.DD_SOURCE_API_URL,
@@ -68,6 +75,13 @@ def handle_parse_result(self, ctx: Context, opts: Dict[Any, Any], args: List[Any
         help="Datadog destination organization APP key.",
         cls=CustomOptionClass,
     ),
+    option(
+        "--destination-jwt",
+        envvar=constants.DD_DESTINATION_JWT,
+        required=False,
+        help="Datadog destination organization JWT (takes precedence over API key).",
+        cls=CustomOptionClass,
+    ),
     option(
         "--destination-api-url",
         envvar=constants.DD_DESTINATION_API_URL,

datadog_sync/constants.py

@@ -9,9 +9,11 @@
 DD_SOURCE_API_URL = "DD_SOURCE_API_URL"
 DD_SOURCE_API_KEY = "DD_SOURCE_API_KEY"
 DD_SOURCE_APP_KEY = "DD_SOURCE_APP_KEY"
+DD_SOURCE_JWT = "DD_SOURCE_JWT"
 DD_DESTINATION_API_URL = "DD_DESTINATION_API_URL"
 DD_DESTINATION_API_KEY = "DD_DESTINATION_API_KEY"
 DD_DESTINATION_APP_KEY = "DD_DESTINATION_APP_KEY"
+DD_DESTINATION_JWT = "DD_DESTINATION_JWT"
 DD_HTTP_CLIENT_RETRY_TIMEOUT = "DD_HTTP_CLIENT_RETRY_TIMEOUT"
 DD_HTTP_CLIENT_TIMEOUT = "DD_HTTP_CLIENT_TIMEOUT"
 DD_RESOURCES = "DD_RESOURCES"

datadog_sync/utils/configuration.py

@@ -144,18 +144,30 @@ def build_config(cmd: Command, **kwargs: Optional[Any]) -> Configuration:
     send_metrics = kwargs.get("send_metrics")
     verify_ssl = kwargs.get("verify_ssl_certificates", True)
 
+    # JWT takes precedence over API keys, so warn if user provided both
+    if (kwargs.get("source_jwt") and kwargs.get("source_api_key")) or (
+        kwargs.get("destination_jwt") and kwargs.get("destination_api_key")
+    ):
+        logger.warning("Both a JWT and an API key were found, the JWT will take precedence.")
+
     source_auth = {}
-    if k := kwargs.get("source_api_key"):
+    # JWT takes precedence over API keys
+    if jwt := kwargs.get("source_jwt"):
+        source_auth["jwtAuth"] = jwt
+    elif k := kwargs.get("source_api_key"):
         source_auth["apiKeyAuth"] = k
-    if k := kwargs.get("source_app_key"):
-        source_auth["appKeyAuth"] = k
+        if k := kwargs.get("source_app_key"):
+            source_auth["appKeyAuth"] = k
     source_client = CustomClient(source_api_url, source_auth, retry_timeout, timeout, send_metrics, verify_ssl)
 
     destination_auth = {}
-    if k := kwargs.get("destination_api_key"):
+    # JWT takes precedence over API keys
+    if jwt := kwargs.get("destination_jwt"):
+        destination_auth["jwtAuth"] = jwt
+    elif k := kwargs.get("destination_api_key"):
         destination_auth["apiKeyAuth"] = k
-    if k := kwargs.get("destination_app_key"):
-        destination_auth["appKeyAuth"] = k
+        if k := kwargs.get("destination_app_key"):
+            destination_auth["appKeyAuth"] = k
     destination_client = CustomClient(
         destination_api_url,
         destination_auth,

datadog_sync/utils/custom_client.py

@@ -291,11 +291,17 @@ async def get_ddr_status(self) -> Dict:
 
 def build_default_headers(auth_obj: Dict[str, str]) -> Dict[str, str]:
     headers = {
-        "DD-API-KEY": auth_obj.get("apiKeyAuth", ""),
-        "DD-APPLICATION-KEY": auth_obj.get("appKeyAuth", ""),
         "Content-Type": "application/json",
         "User-Agent": _get_user_agent(),
     }
+
+    # JWT takes precedence over API keys
+    if jwt := auth_obj.get("jwtAuth"):
+        headers["dd-auth-jwt"] = jwt
+    else:
+        headers["DD-API-KEY"] = auth_obj.get("apiKeyAuth", "")
+        headers["DD-APPLICATION-KEY"] = auth_obj.get("appKeyAuth", "")
+
     return headers