diff --git a/datadog_sync/commands/shared/options.py b/datadog_sync/commands/shared/options.py index d7cf83a1..3904f35b 100644 --- a/datadog_sync/commands/shared/options.py +++ b/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, diff --git a/datadog_sync/constants.py b/datadog_sync/constants.py index a750236e..cdd127f8 100644 --- a/datadog_sync/constants.py +++ b/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" diff --git a/datadog_sync/utils/configuration.py b/datadog_sync/utils/configuration.py index 660e7cb2..852085eb 100644 --- a/datadog_sync/utils/configuration.py +++ b/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, diff --git a/datadog_sync/utils/custom_client.py b/datadog_sync/utils/custom_client.py index 201e7f43..76f09073 100644 --- a/datadog_sync/utils/custom_client.py +++ b/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