From cfccff9cdb0de599c04121f25514d218c987eb1d Mon Sep 17 00:00:00 2001 From: Piotr Roszatycki Date: Fri, 14 Feb 2025 20:52:00 +0100 Subject: [PATCH 1/2] Use STS token for AWS authorization --- .../prometheus/prometheus_utils.py | 35 ++++++++++++--- robusta_krr/core/models/config.py | 43 +++++++++++++++---- robusta_krr/main.py | 38 +++++++++++++--- tests/formatters/test_csv_formatter.py | 13 ++++-- 4 files changed, 106 insertions(+), 23 deletions(-) diff --git a/robusta_krr/core/integrations/prometheus/prometheus_utils.py b/robusta_krr/core/integrations/prometheus/prometheus_utils.py index e561af24..a7144afd 100644 --- a/robusta_krr/core/integrations/prometheus/prometheus_utils.py +++ b/robusta_krr/core/integrations/prometheus/prometheus_utils.py @@ -3,7 +3,12 @@ from typing import TYPE_CHECKING import boto3 -from prometrix import AWSPrometheusConfig, CoralogixPrometheusConfig, PrometheusConfig, VictoriaMetricsPrometheusConfig +from prometrix import ( + AWSPrometheusConfig, + CoralogixPrometheusConfig, + PrometheusConfig, + VictoriaMetricsPrometheusConfig, +) from robusta_krr.core.models.config import settings @@ -37,9 +42,26 @@ def generate_prometheus_config( session = boto3.Session(profile_name=settings.eks_managed_prom_profile_name) credentials = session.get_credentials() credentials = credentials.get_frozen_credentials() - region = settings.eks_managed_prom_region if settings.eks_managed_prom_region else session.region_name - access_key = settings.eks_access_key if settings.eks_access_key else credentials.access_key - secret_key = settings.eks_secret_key.get_secret_value() if settings.eks_secret_key else credentials.secret_key + region = ( + settings.eks_managed_prom_region + if settings.eks_managed_prom_region + else session.region_name + ) + access_key = ( + settings.eks_access_key + if settings.eks_access_key + else credentials.access_key + ) + secret_key = ( + settings.eks_secret_key.get_secret_value() + if settings.eks_secret_key + else credentials.secret_key + ) + token = ( + settings.eks_token.get_secret_value() + if settings.eks_token + else credentials.token + ) service_name = settings.eks_service_name if settings.eks_secret_key else "aps" if not region: raise Exception("No eks region specified") @@ -47,13 +69,16 @@ def generate_prometheus_config( return AWSPrometheusConfig( access_key=access_key, secret_access_key=secret_key, + token=token, aws_region=region, service_name=service_name, **baseconfig, ) # coralogix config if settings.coralogix_token: - return CoralogixPrometheusConfig(**baseconfig, prometheus_token=settings.coralogix_token.get_secret_value()) + return CoralogixPrometheusConfig( + **baseconfig, prometheus_token=settings.coralogix_token.get_secret_value() + ) if isinstance(metrics_service, VictoriaMetricsService): return VictoriaMetricsPrometheusConfig(**baseconfig) return PrometheusConfig(**baseconfig) diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py index 32241ed1..9d8c6bfb 100644 --- a/robusta_krr/core/models/config.py +++ b/robusta_krr/core/models/config.py @@ -44,6 +44,7 @@ class Config(pd.BaseSettings): eks_managed_prom_profile_name: Optional[str] = pd.Field(None) eks_access_key: Optional[str] = pd.Field(None) eks_secret_key: Optional[pd.SecretStr] = pd.Field(None) + eks_token: Optional[pd.SecretStr] = pd.Field(None) eks_service_name: Optional[str] = pd.Field(None) eks_managed_prom_region: Optional[str] = pd.Field(None) coralogix_token: Optional[pd.SecretStr] = pd.Field(None) @@ -91,32 +92,46 @@ def validate_prometheus_url(cls, v: Optional[str]): return v @pd.validator("prometheus_other_headers", pre=True) - def validate_prometheus_other_headers(cls, headers: Union[list[str], dict[str, str]]) -> dict[str, str]: + def validate_prometheus_other_headers( + cls, headers: Union[list[str], dict[str, str]] + ) -> dict[str, str]: if isinstance(headers, dict): return headers - return {k.strip().lower(): v.strip() for k, v in [header.split(":") for header in headers]} + return { + k.strip().lower(): v.strip() + for k, v in [header.split(":") for header in headers] + } @pd.validator("namespaces") - def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]: + def validate_namespaces( + cls, v: Union[list[str], Literal["*"]] + ) -> Union[list[str], Literal["*"]]: if v == []: return "*" if isinstance(v, list): for val in v: if val.startswith("*"): - raise ValueError("Namespace's values cannot start with an asterisk (*)") + raise ValueError( + "Namespace's values cannot start with an asterisk (*)" + ) return [val.lower() for val in v] @pd.validator("resources", pre=True) - def validate_resources(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]: + def validate_resources( + cls, v: Union[list[str], Literal["*"]] + ) -> Union[list[str], Literal["*"]]: if v == []: return "*" # NOTE: KindLiteral.__args__ is a tuple of all possible values of KindLiteral # So this will preserve the big and small letters of the resource - return [next(r for r in KindLiteral.__args__ if r.lower() == val.lower()) for val in v] + return [ + next(r for r in KindLiteral.__args__ if r.lower() == val.lower()) + for val in v + ] def create_strategy(self) -> AnyStrategy: StrategyType = AnyStrategy.find(self.strategy) @@ -140,7 +155,9 @@ def context(self) -> Optional[str]: @property def logging_console(self) -> Console: if getattr(self, "_logging_console") is None: - self._logging_console = Console(file=sys.stderr if self.log_to_stderr else sys.stdout, width=self.width) + self._logging_console = Console( + file=sys.stderr if self.log_to_stderr else sys.stdout, width=self.width + ) return self._logging_console def load_kubeconfig(self) -> None: @@ -155,7 +172,9 @@ def get_kube_client(self, context: Optional[str] = None): if context is None: return None - api_client = config.new_client_from_config(context=context, config_file=self.kubeconfig) + api_client = config.new_client_from_config( + context=context, config_file=self.kubeconfig + ) if self.impersonate_user is not None: # trick copied from https://github.com/kubernetes-client/python/issues/362 api_client.set_default_header("Impersonate-User", self.impersonate_user) @@ -175,7 +194,13 @@ def set_config(config: Config) -> None: handlers=[RichHandler(console=config.logging_console)], ) logging.getLogger("").setLevel(logging.CRITICAL) - logger.setLevel(logging.DEBUG if config.verbose else logging.CRITICAL if config.quiet else logging.INFO) + logger.setLevel( + logging.DEBUG + if config.verbose + else logging.CRITICAL + if config.quiet + else logging.INFO + ) @staticmethod def get_config() -> Optional[Config]: diff --git a/robusta_krr/main.py b/robusta_krr/main.py index 7159bdde..d88e09b1 100644 --- a/robusta_krr/main.py +++ b/robusta_krr/main.py @@ -171,6 +171,12 @@ def run_strategy( help="Sets the secret key for eks prometheus connection.", rich_help_panel="Prometheus EKS Settings", ), + eks_token: Optional[str] = typer.Option( + None, + "--eks-token", + help="Sets the session token for eks prometheus connection.", + rich_help_panel="Prometheus EKS Settings", + ), eks_service_name: Optional[str] = typer.Option( "aps", "--eks-service-name", @@ -234,13 +240,24 @@ def run_strategy( rich_help_panel="Output Settings", ), verbose: bool = typer.Option( - False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings" + False, + "--verbose", + "-v", + help="Enable verbose mode", + rich_help_panel="Logging Settings", ), quiet: bool = typer.Option( - False, "--quiet", "-q", help="Enable quiet mode", rich_help_panel="Logging Settings" + False, + "--quiet", + "-q", + help="Enable quiet mode", + rich_help_panel="Logging Settings", ), log_to_stderr: bool = typer.Option( - False, "--logtostderr", help="Pass logs to stderr", rich_help_panel="Logging Settings" + False, + "--logtostderr", + help="Pass logs to stderr", + rich_help_panel="Logging Settings", ), width: Optional[int] = typer.Option( None, @@ -270,7 +287,10 @@ def run_strategy( ) -> None: f"""Run KRR using the `{_strategy_name}` strategy""" if not show_severity and format != "csv": - raise click.BadOptionUsage("--exclude-severity", "--exclude-severity works only with format=csv") + raise click.BadOptionUsage( + "--exclude-severity", + "--exclude-severity works only with format=csv", + ) try: config = Config( @@ -292,6 +312,7 @@ def run_strategy( eks_managed_prom_profile_name=eks_managed_prom_profile_name, eks_access_key=eks_access_key, eks_secret_key=eks_secret_key, + eks_token=eks_token, eks_service_name=eks_service_name, coralogix_token=coralogix_token, openshift=openshift, @@ -329,7 +350,14 @@ def run_strategy( kind=inspect.Parameter.KEYWORD_ONLY, default=OptionInfo( default=field_meta.default, - param_decls=list(set([f"--{field_name}", f"--{field_name.replace('_', '-')}"])), + param_decls=list( + set( + [ + f"--{field_name}", + f"--{field_name.replace('_', '-')}", + ] + ) + ), help=f"{field_meta.field_info.description}", rich_help_panel="Strategy Settings", ), diff --git a/tests/formatters/test_csv_formatter.py b/tests/formatters/test_csv_formatter.py index 150b2aa5..95a77947 100644 --- a/tests/formatters/test_csv_formatter.py +++ b/tests/formatters/test_csv_formatter.py @@ -119,6 +119,7 @@ "eks_managed_prom_profile_name": null, "eks_access_key": null, "eks_secret_key": null, + "eks_token": null, "eks_service_name": "aps", "eks_managed_prom_region": null, "coralogix_token": null, @@ -215,7 +216,9 @@ def _load_result(override_config: dict[str, Any]) -> Result: ), ], ) -def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str]) -> None: +def test_csv_headers( + override_config: dict[str, Any], expected_headers: list[str] +) -> None: result = _load_result(override_config=override_config) output = csv_exporter(result) reader = csv.DictReader(io.StringIO(output)) @@ -235,7 +238,7 @@ def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str "Old Pods": "1", "Type": "Deployment", "Container": "mock-container-1", - 'Severity': 'CRITICAL', + "Severity": "CRITICAL", "CPU Diff": "-87m", "CPU Requests": "(-43m) 50m -> 6m", "CPU Limits": "2.0 -> ?", @@ -271,7 +274,7 @@ def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str "Old Pods": "1", "Type": "Deployment", "Container": "mock-container-1", - 'Severity': 'CRITICAL', + "Severity": "CRITICAL", "CPU Diff": "-87m", "CPU Requests": "(-43m) 50m -> 6m", "CPU Limits": "2.0 -> ?", @@ -282,7 +285,9 @@ def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str ), ], ) -def test_csv_row_value(override_config: dict[str, Any], expected_first_row: list[str]) -> None: +def test_csv_row_value( + override_config: dict[str, Any], expected_first_row: list[str] +) -> None: result = _load_result(override_config=override_config) output = csv_exporter(result) reader = csv.DictReader(io.StringIO(output)) From 6b2c5b8e5f81e1cc1c736485c92a7a44fee5e239 Mon Sep 17 00:00:00 2001 From: Piotr Roszatycki Date: Fri, 14 Feb 2025 21:00:13 +0100 Subject: [PATCH 2/2] Black --- .../prometheus/prometheus_utils.py | 28 +++---------- robusta_krr/core/models/config.py | 42 ++++--------------- tests/formatters/test_csv_formatter.py | 8 +--- 3 files changed, 16 insertions(+), 62 deletions(-) diff --git a/robusta_krr/core/integrations/prometheus/prometheus_utils.py b/robusta_krr/core/integrations/prometheus/prometheus_utils.py index a7144afd..ec6f20e2 100644 --- a/robusta_krr/core/integrations/prometheus/prometheus_utils.py +++ b/robusta_krr/core/integrations/prometheus/prometheus_utils.py @@ -42,26 +42,10 @@ def generate_prometheus_config( session = boto3.Session(profile_name=settings.eks_managed_prom_profile_name) credentials = session.get_credentials() credentials = credentials.get_frozen_credentials() - region = ( - settings.eks_managed_prom_region - if settings.eks_managed_prom_region - else session.region_name - ) - access_key = ( - settings.eks_access_key - if settings.eks_access_key - else credentials.access_key - ) - secret_key = ( - settings.eks_secret_key.get_secret_value() - if settings.eks_secret_key - else credentials.secret_key - ) - token = ( - settings.eks_token.get_secret_value() - if settings.eks_token - else credentials.token - ) + region = settings.eks_managed_prom_region if settings.eks_managed_prom_region else session.region_name + access_key = settings.eks_access_key if settings.eks_access_key else credentials.access_key + secret_key = settings.eks_secret_key.get_secret_value() if settings.eks_secret_key else credentials.secret_key + token = settings.eks_token.get_secret_value() if settings.eks_token else credentials.token service_name = settings.eks_service_name if settings.eks_secret_key else "aps" if not region: raise Exception("No eks region specified") @@ -76,9 +60,7 @@ def generate_prometheus_config( ) # coralogix config if settings.coralogix_token: - return CoralogixPrometheusConfig( - **baseconfig, prometheus_token=settings.coralogix_token.get_secret_value() - ) + return CoralogixPrometheusConfig(**baseconfig, prometheus_token=settings.coralogix_token.get_secret_value()) if isinstance(metrics_service, VictoriaMetricsService): return VictoriaMetricsPrometheusConfig(**baseconfig) return PrometheusConfig(**baseconfig) diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py index 9d8c6bfb..5735387f 100644 --- a/robusta_krr/core/models/config.py +++ b/robusta_krr/core/models/config.py @@ -92,46 +92,32 @@ def validate_prometheus_url(cls, v: Optional[str]): return v @pd.validator("prometheus_other_headers", pre=True) - def validate_prometheus_other_headers( - cls, headers: Union[list[str], dict[str, str]] - ) -> dict[str, str]: + def validate_prometheus_other_headers(cls, headers: Union[list[str], dict[str, str]]) -> dict[str, str]: if isinstance(headers, dict): return headers - return { - k.strip().lower(): v.strip() - for k, v in [header.split(":") for header in headers] - } + return {k.strip().lower(): v.strip() for k, v in [header.split(":") for header in headers]} @pd.validator("namespaces") - def validate_namespaces( - cls, v: Union[list[str], Literal["*"]] - ) -> Union[list[str], Literal["*"]]: + def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]: if v == []: return "*" if isinstance(v, list): for val in v: if val.startswith("*"): - raise ValueError( - "Namespace's values cannot start with an asterisk (*)" - ) + raise ValueError("Namespace's values cannot start with an asterisk (*)") return [val.lower() for val in v] @pd.validator("resources", pre=True) - def validate_resources( - cls, v: Union[list[str], Literal["*"]] - ) -> Union[list[str], Literal["*"]]: + def validate_resources(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]: if v == []: return "*" # NOTE: KindLiteral.__args__ is a tuple of all possible values of KindLiteral # So this will preserve the big and small letters of the resource - return [ - next(r for r in KindLiteral.__args__ if r.lower() == val.lower()) - for val in v - ] + return [next(r for r in KindLiteral.__args__ if r.lower() == val.lower()) for val in v] def create_strategy(self) -> AnyStrategy: StrategyType = AnyStrategy.find(self.strategy) @@ -155,9 +141,7 @@ def context(self) -> Optional[str]: @property def logging_console(self) -> Console: if getattr(self, "_logging_console") is None: - self._logging_console = Console( - file=sys.stderr if self.log_to_stderr else sys.stdout, width=self.width - ) + self._logging_console = Console(file=sys.stderr if self.log_to_stderr else sys.stdout, width=self.width) return self._logging_console def load_kubeconfig(self) -> None: @@ -172,9 +156,7 @@ def get_kube_client(self, context: Optional[str] = None): if context is None: return None - api_client = config.new_client_from_config( - context=context, config_file=self.kubeconfig - ) + api_client = config.new_client_from_config(context=context, config_file=self.kubeconfig) if self.impersonate_user is not None: # trick copied from https://github.com/kubernetes-client/python/issues/362 api_client.set_default_header("Impersonate-User", self.impersonate_user) @@ -194,13 +176,7 @@ def set_config(config: Config) -> None: handlers=[RichHandler(console=config.logging_console)], ) logging.getLogger("").setLevel(logging.CRITICAL) - logger.setLevel( - logging.DEBUG - if config.verbose - else logging.CRITICAL - if config.quiet - else logging.INFO - ) + logger.setLevel(logging.DEBUG if config.verbose else logging.CRITICAL if config.quiet else logging.INFO) @staticmethod def get_config() -> Optional[Config]: diff --git a/tests/formatters/test_csv_formatter.py b/tests/formatters/test_csv_formatter.py index 95a77947..98cffbcf 100644 --- a/tests/formatters/test_csv_formatter.py +++ b/tests/formatters/test_csv_formatter.py @@ -216,9 +216,7 @@ def _load_result(override_config: dict[str, Any]) -> Result: ), ], ) -def test_csv_headers( - override_config: dict[str, Any], expected_headers: list[str] -) -> None: +def test_csv_headers(override_config: dict[str, Any], expected_headers: list[str]) -> None: result = _load_result(override_config=override_config) output = csv_exporter(result) reader = csv.DictReader(io.StringIO(output)) @@ -285,9 +283,7 @@ def test_csv_headers( ), ], ) -def test_csv_row_value( - override_config: dict[str, Any], expected_first_row: list[str] -) -> None: +def test_csv_row_value(override_config: dict[str, Any], expected_first_row: list[str]) -> None: result = _load_result(override_config=override_config) output = csv_exporter(result) reader = csv.DictReader(io.StringIO(output))