LLM/pentestagent
1
0
Fork 0
mirror of https://github.com/GH05TCREW/pentestagent.git synced 2026-03-08 06:44:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

style: resolve linting errors and improve formatting

Browse Source
This commit is contained in:
GH05TCREW
2026-01-11 16:00:43 -07:00
parent cb7597e74b
commit 49c0b93b59
4 changed files with 37 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
.env.example
View File

@@ -15,11 +15,6 @@ TAVILY_API_KEY=
# Other providers: azure/, bedrock/, groq/, ollama/, together_ai/ (see litellm docs)
PENTESTAGENT_MODEL=gpt-5
# Provider selection:
# Note: The app determines provider from `PENTESTAGENT_MODEL` prefix
# (e.g., `ollama/...`, `gpt-5`, `claude-...`, `gemini/...`). No separate
# `LLM_PROVIDER` variable is used.
# Ollama base URL (set this when using an `ollama/...` model)
# Example: http://127.0.0.1:11434 or http://192.168.0.165:11434
OLLAMA_BASE_URL=http://127.0.0.1:11434
@@ -32,7 +27,7 @@ OLLAMA_BASE_URL=http://127.0.0.1:11434
PENTESTAGENT_EMBEDDINGS=local
# Settings
PENTESTAGENT_DEBUG=true
PENTESTAGENT_DEBUG=false
# Optional: manually declare model/context and daily token budgeting
# Useful when provider metadata isn't available or you want to enforce local limits.

8
pentestagent/interface/cli.py
View File

@@ -1,5 +1,6 @@
"""Non-interactive CLI mode for PentestAgent."""
import ast
import asyncio
import time
from datetime import datetime
@@ -9,9 +10,6 @@ from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.text import Text
import json
import ast
import re
from ..config.constants import AGENT_MAX_ITERATIONS, ORCHESTRATOR_MAX_ITERATIONS
@@ -206,7 +204,9 @@ async def run_cli(
except Exception:
parsed_for_summary = None
if isinstance(parsed_for_summary, dict) and isinstance(parsed_for_summary.get("summary"), str):
if isinstance(parsed_for_summary, dict) and isinstance(
parsed_for_summary.get("summary"), str
):
console.print(
Panel(
Markdown(parsed_for_summary.get("summary")),

41
pentestagent/interface/tui.py
View File

@@ -426,13 +426,15 @@ class MemoryDiagnostics(Static):
# Use a consistent bar width for all bars and align labels
bar_width = 28
labels = ["Tokens:", "Messages:", "Retention:"]
label_width = max(len(l) for l in labels)
label_width = max(len(label_text) for label_text in labels)
# Tokens line
ratio = current_tokens / max(1, budget)
bar = self._bar(ratio, width=bar_width)
label = "Tokens:".ljust(label_width)
text.append(f"{label} [{bar}] {current_tokens:,} / {budget:,}\n", style="#9a9a9a")
text.append(
f"{label} [{bar}] {current_tokens:,} / {budget:,}\n", style="#9a9a9a"
)
# Messages line (scale messages to an expected max window)
expected_msgs_max = max(1, recent_keep * 6)
@@ -445,7 +447,9 @@ class MemoryDiagnostics(Static):
k_ratio = min(1.0, recent_keep / max(1, recent_keep))
keep_bar = self._bar(k_ratio, width=bar_width)
label = "Retention:".ljust(label_width)
text.append(f"{label} [{keep_bar}] keeping last {recent_keep}\n", style="#9a9a9a")
text.append(
f"{label} [{keep_bar}] keeping last {recent_keep}\n", style="#9a9a9a"
)
# Summary status
summary_state = "active" if has_summary else "inactive"
@@ -453,7 +457,9 @@ class MemoryDiagnostics(Static):
text.append(f"Summary: {emoji} {summary_state}\n", style="#9a9a9a")
# Summarized / threshold
text.append(f"Summarized: {summarized_count} / {thresh:,}\n", style="#9a9a9a")
text.append(
f"Summarized: {summarized_count} / {thresh:,}\n", style="#9a9a9a"
)
text.append(f"Threshold: {thresh:,}\n", style="#9a9a9a")
except Exception as e:
@@ -501,7 +507,10 @@ class TokenDiagnostics(Static):
text.append("Token Usage Diagnostics\n", style="bold #d4d4d4")
if not token_tracker:
text.append("Token tracker not available (tools/token_tracker).\n", style="#9a9a9a")
text.append(
"Token tracker not available (tools/token_tracker).\n",
style="#9a9a9a",
)
return text
stats = token_tracker.get_stats_sync()
@@ -569,7 +578,10 @@ class TokenDiagnostics(Static):
text.append("Environment configuration errors:\n", style="#ef4444")
for e in env_errors:
text.append(f" - {e}\n", style="#9a9a9a")
text.append("\nSet environment variables correctly to compute costs.\n", style="#9a9a9a")
text.append(
"\nSet environment variables correctly to compute costs.\n",
style="#9a9a9a",
)
return text
# Compute costs
@@ -580,7 +592,10 @@ class TokenDiagnostics(Static):
else:
# Require per-direction costs to be present to compute
if input_cost_per_m is None or output_cost_per_m is None:
text.append("Cost vars missing. Set COST_PER_MILLION or both INPUT_COST_PER_MILLION and OUTPUT_COST_PER_MILLION.\n", style="#9a9a9a")
text.append(
"Cost vars missing. Set COST_PER_MILLION or both INPUT_COST_PER_MILLION and OUTPUT_COST_PER_MILLION.\n",
style="#9a9a9a",
)
# Still show numeric token stats below
input_cost = output_cost = None
else:
@@ -618,7 +633,7 @@ class TokenDiagnostics(Static):
"Current date:",
"Reset occurred:",
]
label_width = max(len(l) for l in labels)
label_width = max(len(label_text) for label_text in labels)
# Last command tokens
label = "Last command:".ljust(label_width)
@@ -1137,6 +1152,7 @@ class PentestAgentTUI(App):
# Mount a new diagnostics panel with a unique ID and scroll into view
try:
import uuid
panel_id = f"memory-diagnostics-{uuid.uuid4().hex}"
except Exception:
panel_id = None
@@ -1158,6 +1174,7 @@ class PentestAgentTUI(App):
# Mount a new diagnostics panel with a unique ID and scroll into view
try:
import uuid
panel_id = f"token-diagnostics-{uuid.uuid4().hex}"
except Exception:
panel_id = None
@@ -1275,10 +1292,14 @@ class PentestAgentTUI(App):
scroll = self.query_one("#chat-scroll", ScrollableContainer)
updated = False
for child in scroll.children:
if isinstance(child, SystemMessage) and "PentestAgent ready" in getattr(child, "message_content", ""):
if isinstance(child, SystemMessage) and "PentestAgent ready" in getattr(
child, "message_content", ""
):
# Append Target line if not already present
if "Target:" not in child.message_content:
child.message_content = child.message_content + f"\n Target: {target}"
child.message_content = (
child.message_content + f"\n Target: {target}"
)
try:
child.refresh()
except Exception:

2
pentestagent/tools/token_tracker.py
View File

@@ -9,7 +9,7 @@ import json
import threading
from datetime import date
from pathlib import Path
from typing import Dict, Any
from typing import Any, Dict
# Persistent storage (loot root)
_data_file: Path = Path("loot/token_usage.json")