Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions robusta_krr/core/models/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,11 @@ class Config(pd.BaseSettings):
width: Optional[int] = pd.Field(None, ge=1)
show_severity: bool = True

# Publishing to url settings
publish_scan_url: Optional[str] = pd.Field(None)
start_time: Optional[str] = pd.Field(None)
scan_id: Optional[str] = pd.Field(None)

# Output Settings
file_output: Optional[str] = pd.Field(None)
file_output_dynamic: bool = pd.Field(False)
Expand Down
71 changes: 70 additions & 1 deletion robusta_krr/core/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,18 @@
import math
import os
import sys
import time
import warnings
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Union
from datetime import timedelta, datetime
from prometrix import PrometheusNotFound
from rich.console import Console
from slack_sdk import WebClient

from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import requests
import json
import traceback
from robusta_krr.core.abstract.strategies import ResourceRecommendation, RunResult
from robusta_krr.core.integrations.kubernetes import KubernetesLoader
from robusta_krr.core.integrations.prometheus import ClusterNotSpecifiedException, PrometheusMetricsLoader
Expand Down Expand Up @@ -104,6 +108,8 @@ def _process_result(self, result: Result) -> None:
result.errors = self.errors

Formatter = settings.Formatter

self._send_result(settings.publish_scan_url, settings.start_time, settings.scan_id, result)
formatted = result.format(Formatter)
rich = getattr(Formatter, "__rich_console__", False)

Expand Down Expand Up @@ -329,6 +335,7 @@ async def run(self) -> int:
except Exception as e:
logger.error(f"Could not load kubernetes configuration: {e}")
logger.error("Try to explicitly set --context and/or --kubeconfig flags.")
publish_error(f"Could not load kubernetes configuration: {e}")
return 1 # Exit with error

try:
Expand All @@ -348,9 +355,71 @@ async def run(self) -> int:
self._process_result(result)
except (ClusterNotSpecifiedException, CriticalRunnerException) as e:
logger.critical(e)
publish_error(traceback.format_exc())
return 1 # Exit with error
except Exception:
logger.exception("An unexpected error occurred")
publish_error(traceback.format_exc())
return 1 # Exit with error
else:
return 0 # Exit with success

def _send_result(self, url: str, start_time: datetime, scan_id: str, result: Result):
result_dict = json.loads(result.json(indent=2))
_send_scan_payload(url, scan_id, start_time, result_dict, is_error=False)

def publish_input_error(url: str, scan_id: str, start_time: str, error: str):
_send_scan_payload(url, scan_id, start_time, error, is_error=True)

def publish_error(error: str):
_send_scan_payload(settings.publish_scan_url, settings.scan_id, settings.start_time, error, is_error=True)

@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=1, max=8),
reraise=True
)
def _post_scan_request(url: str, headers: dict, payload: dict, scan_id: str, is_error: bool):
logger_msg = "Sending error scan" if is_error else "Sending scan"
logger.info(f"{logger_msg} for scan_id={scan_id} to url={url}")
response = requests.post(url, headers=headers, json=payload)
logger.info(f"scan_id={scan_id} | Status code: {response.status_code}")
logger.info(f"scan_id={scan_id} | Response body: {response.text}")
return response


def _send_scan_payload(
url: str,
scan_id: str,
start_time: Union[str, datetime],
result_data: Union[str, dict],
is_error: bool = False
):
if not url or not scan_id or not start_time:
logger.debug(f"Missing required parameters: url={bool(url)}, scan_id={bool(scan_id)}, start_time={bool(start_time)}")
return

logger.debug(f"Preparing to send scan payload. scan_id={scan_id}, is_error={is_error}")

headers = {"Content-Type": "application/json"}

if isinstance(start_time, datetime):
logger.debug(f"Converting datetime to ISO format for scan_id={scan_id}")
start_time = start_time.isoformat()

action_request = {
"action_name": "process_scan",
"action_params": {
"result": result_data,
"scan_type": "krr",
"scan_id": scan_id,
"start_time": start_time,
}
}

try:
_post_scan_request(url, headers, action_request, scan_id, is_error)
except requests.exceptions.RequestException as e:
logger.error(f"scan_id={scan_id} | All retry attempts failed due to RequestException: {e}", exc_info=True)
except Exception as e:
logger.error(f"scan_id={scan_id} | Unexpected error after retries: {e}", exc_info=True)
28 changes: 25 additions & 3 deletions robusta_krr/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from robusta_krr.core.abstract import formatters
from robusta_krr.core.abstract.strategies import BaseStrategy
from robusta_krr.core.models.config import Config
from robusta_krr.core.runner import Runner
from robusta_krr.core.runner import Runner, publish_input_error
from robusta_krr.utils.version import get_version

app = typer.Typer(
Expand Down Expand Up @@ -266,6 +266,24 @@ def run_strategy(
help="Send to output to a slack channel, must have SLACK_BOT_TOKEN",
rich_help_panel="Output Settings",
),
publish_scan_url: Optional[str] = typer.Option(
None,
"--publish_scan_url",
help="Sends the output to a robusta_runner instance",
rich_help_panel="Publish Scan Settings",
),
start_time: Optional[str] = typer.Option(
None,
"--start_time",
help="Start time of the scan",
rich_help_panel="Publish Scan Settings",
),
scan_id: Optional[str] = typer.Option(
None,
"--scan_id",
help="A UUID scan identifier",
rich_help_panel="Publish Scan Settings",
),
**strategy_args,
) -> None:
f"""Run KRR using the `{_strategy_name}` strategy"""
Expand Down Expand Up @@ -310,10 +328,14 @@ def run_strategy(
show_severity=show_severity,
strategy=_strategy_name,
other_args=strategy_args,
)
publish_scan_url=publish_scan_url,
start_time=start_time,
scan_id=scan_id,
)
Config.set_config(config)
except ValidationError:
except ValidationError as e:
logger.exception("Error occured while parsing arguments")
publish_input_error( publish_scan_url, start_time, scan_id, str(e))
else:
runner = Runner()
exit_code = asyncio.run(runner.run())
Expand Down