diff --git a/agentic_security/report_chart.py b/agentic_security/report_chart.py index 45387cb..e1cd4c4 100644 --- a/agentic_security/report_chart.py +++ b/agentic_security/report_chart.py @@ -7,11 +7,31 @@ import pandas as pd from matplotlib.cm import ScalarMappable from matplotlib.colors import LinearSegmentedColormap, Normalize +from agentic_security.logutils import logger + from .primitives import Table def plot_security_report(table: Table) -> io.BytesIO: + try: + return _plot_security_report(table=table) + except (TypeError, ValueError, OverflowError, IndexError, Exception) as e: + logger.error(f"Error in generating the security report: {e}") + return io.BytesIO() + + +def generate_identifiers(data: pd.DataFrame) -> list[str]: + try: + _generate_identifiers(data=data) + except (TypeError, ValueError, Exception) as e: + logger.error(f"Error in generate_identifiers: {e}") + return [""] + + +def _plot_security_report(table: Table) -> io.BytesIO: # Data preprocessing + logger.info("Data preprocessing started.") + data = pd.DataFrame(table) # Sort by failure rate and reset index @@ -22,10 +42,10 @@ def plot_security_report(table: Table) -> io.BytesIO: fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"}) fig.set_facecolor("#f0f0f0") ax.set_facecolor("#f0f0f0") + logger.info("Plot setup complete.") # Styling parameters colors = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"][::-1] # Pastel palette - # colors = ["#440154", "#3b528b", "#21908c", "#5dc863"] # Viridis-inspired palette cmap = LinearSegmentedColormap.from_list("custom", colors, N=256) norm = Normalize(vmin=data["tokens"].min(), vmax=data["tokens"].max()) @@ -76,7 +96,10 @@ def plot_security_report(table: Table) -> io.BytesIO: # Title and caption fig.suptitle( - "Security Report for Different Modules", fontsize=16, fontweight="bold", y=1.02 + "Security Report for Different Modules", + fontsize=16, + fontweight="bold", + y=1.02, ) caption = "Report generated by https://github.com/msoedov/agentic_security" fig.text( @@ -114,17 +137,12 @@ def plot_security_report(table: Table) -> io.BytesIO: data["identifier"], data["failureRate"], data["module"] ) ] - table = ax.table( - cellText=table_data, - loc="right", - cellLoc="left", - ) + table = ax.table(cellText=table_data, loc="right", cellLoc="left") table.auto_set_font_size(False) table.set_fontsize(8) # Adjust table style table.scale(1, 0.7) - for (row, col), cell in table.get_celld().items(): cell.set_edgecolor("none") cell.set_facecolor("#f0f0f0" if row % 2 == 0 else "#e0e0e0") @@ -134,17 +152,18 @@ def plot_security_report(table: Table) -> io.BytesIO: cell.set_text_props(fontweight="bold") # Adjust layout and save - plt.tight_layout() buf = io.BytesIO() plt.savefig(buf, format="png", dpi=300, bbox_inches="tight") plt.close(fig) buf.seek(0) + logger.info("Report successfully generated and saved to buffer.") return buf -def generate_identifiers(data: pd.DataFrame) -> list[str]: +def _generate_identifiers(data: pd.DataFrame) -> list[str]: data_length = len(data) + alphabet = string.ascii_uppercase num_letters = len(alphabet)