Coverage for src/alprina_cli/agents/llm_enhancer.py: 35%
77 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-14 11:27 +0100
1"""
2LLM Enhancement Layer - Claude AI Integration
3Enhances static analysis findings with contextual reasoning
4"""
6import os
7import re
8from typing import Dict, Any, Optional
9from dataclasses import dataclass, asdict
10from loguru import logger
12try:
13 from anthropic import Anthropic
14 ANTHROPIC_AVAILABLE = True
15except ImportError:
16 ANTHROPIC_AVAILABLE = False
17 logger.warning("Anthropic SDK not installed. LLM enhancement unavailable.")
20@dataclass
21class EnhancedVulnerability:
22 """Vulnerability enhanced with LLM reasoning"""
23 # Original fields
24 vulnerability_type: str
25 severity: str
26 title: str
27 description: str
28 file_path: str
29 line_number: int
31 # LLM enhancements
32 business_impact: str = ""
33 economic_loss: str = ""
34 attack_scenario: str = ""
35 historical_precedent: str = ""
36 priority_reasoning: str = ""
37 remediation_code: str = ""
38 llm_explanation: str = ""
39 llm_enhanced: bool = False
41 def to_dict(self) -> Dict[str, Any]:
42 """Convert to dictionary for JSON serialization"""
43 return asdict(self)
46class LLMEnhancer:
47 """Enhance security findings with Claude AI"""
49 def __init__(self, api_key: Optional[str] = None):
50 """Initialize Claude client
52 Args:
53 api_key: Optional API key, defaults to ANTHROPIC_API_KEY env var
55 Raises:
56 ValueError: If API key not provided and not in environment
57 ImportError: If anthropic package not installed
58 """
59 if not ANTHROPIC_AVAILABLE:
60 raise ImportError(
61 "Anthropic SDK not installed. Install with: pip install anthropic"
62 )
64 api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
65 if not api_key:
66 raise ValueError(
67 "ANTHROPIC_API_KEY environment variable not set. "
68 "Get your API key at https://console.anthropic.com/"
69 )
71 self.client = Anthropic(api_key=api_key)
72 self.model = "claude-sonnet-4-5-20250929" # Claude Sonnet 4.5
73 logger.info("✅ LLM Enhancer initialized with Claude Sonnet 4.5")
75 def enhance_vulnerability(
76 self,
77 vuln: Dict[str, Any],
78 contract_code: str
79 ) -> EnhancedVulnerability:
80 """
81 Enhance vulnerability with AI-powered context
83 Args:
84 vuln: Static analysis vulnerability dict
85 contract_code: Full contract source code
87 Returns:
88 EnhancedVulnerability with LLM analysis
89 """
90 try:
91 # Get code context
92 code_context = self._get_code_context(
93 contract_code,
94 vuln.get('line_number', 1)
95 )
97 # Build prompt
98 prompt = self._build_prompt(vuln, code_context)
100 # Call Claude
101 logger.debug(f"Enhancing {vuln.get('title', 'Unknown')} with Claude AI")
102 response = self.client.messages.create(
103 model=self.model,
104 max_tokens=2000,
105 temperature=0, # Deterministic for security
106 system="You are an expert smart contract security auditor with deep knowledge of historical exploits and attack patterns.",
107 messages=[{"role": "user", "content": prompt}]
108 )
110 analysis = response.content[0].text
112 # Parse response
113 enhanced = EnhancedVulnerability(
114 vulnerability_type=vuln.get('vulnerability_type', 'unknown'),
115 severity=vuln.get('severity', 'medium'),
116 title=vuln.get('title', 'Untitled'),
117 description=vuln.get('description', ''),
118 file_path=vuln.get('file_path', ''),
119 line_number=vuln.get('line_number', 0),
120 llm_explanation=analysis,
121 business_impact=self._extract_section(analysis, "Business Impact"),
122 economic_loss=self._extract_section(analysis, "Economic Loss"),
123 attack_scenario=self._extract_section(analysis, "Attack Scenario"),
124 historical_precedent=self._extract_section(analysis, "Historical Precedent"),
125 priority_reasoning=self._extract_section(analysis, "Priority"),
126 remediation_code=self._extract_code_block(analysis),
127 llm_enhanced=True
128 )
130 logger.debug(f"✅ Enhanced vulnerability: {vuln.get('title', 'Unknown')}")
131 return enhanced
133 except Exception as e:
134 logger.error(f"LLM enhancement failed: {e}")
135 # Return original vuln without enhancement
136 return EnhancedVulnerability(
137 vulnerability_type=vuln.get('vulnerability_type', 'unknown'),
138 severity=vuln.get('severity', 'medium'),
139 title=vuln.get('title', 'Untitled'),
140 description=vuln.get('description', ''),
141 file_path=vuln.get('file_path', ''),
142 line_number=vuln.get('line_number', 0),
143 llm_enhanced=False
144 )
146 def _build_prompt(self, vuln: Dict, code_context: str) -> str:
147 """Build Claude prompt for vulnerability analysis"""
148 return f"""Analyze this smart contract vulnerability:
150STATIC ANALYSIS FINDING:
151- Type: {vuln.get('vulnerability_type', 'Unknown')}
152- Severity: {vuln.get('severity', 'Unknown')}
153- Title: {vuln.get('title', 'Untitled')}
154- Description: {vuln.get('description', 'No description')}
155- Location: Line {vuln.get('line_number', 'Unknown')}
157CODE CONTEXT:
158```solidity
159{code_context}
160```
162Provide comprehensive analysis in these sections:
164## Business Impact
165Explain in simple terms how this vulnerability could be exploited and what the consequences would be. Focus on the "what happens" not the technical details.
167## Economic Loss
168Estimate potential financial loss based on 2024 DeFi exploit data:
169- Reentrancy attacks: $35.7M average
170- Access Control issues: $953M total (2024)
171- Oracle Manipulation: $8.7M average
172- Flash Loan attacks: $33.8M average
174## Attack Scenario
175Provide step-by-step technical breakdown of how an attacker would exploit this. Be specific about:
1761. What the attacker does
1772. What happens in the contract
1783. How they profit
1794. How long it takes
181## Historical Precedent
182Reference real exploits with similar patterns. Examples:
183- DAO Hack (2016) - $60M reentrancy
184- Cream Finance (2021) - $130M flash loan + reentrancy
185- Polter Finance (2024) - $8.7M oracle manipulation
186- Wormhole (2022) - $325M signature verification
188## Priority
189Explain why this should (or shouldn't) be fixed immediately. Consider:
190- Exploitability: How easy is it to exploit?
191- Impact: How much damage could it cause?
192- Likelihood: Is the contract deployed? What's the TVL?
194## Remediation Code
195Provide secure code that fixes the vulnerability. Include:
196- Complete function with fix
197- Security comments explaining what changed
198- Any additional imports needed (e.g., OpenZeppelin)
200Format with proper Solidity syntax."""
202 def _get_code_context(self, code: str, line_num: int, context_lines: int = 5) -> str:
203 """Extract code context around vulnerability line"""
204 lines = code.split('\n')
205 start = max(0, line_num - context_lines - 1)
206 end = min(len(lines), line_num + context_lines)
208 context = []
209 for i, line in enumerate(lines[start:end], start=start):
210 marker = " --> " if i == line_num - 1 else " "
211 context.append(f"{marker}{i+1:4d} | {line}")
213 return '\n'.join(context)
215 def _extract_section(self, text: str, section_name: str) -> str:
216 """Extract named section from LLM response"""
217 pattern = rf"## {section_name}\s*\n(.*?)(?=\n##|\Z)"
218 match = re.search(pattern, text, re.DOTALL)
219 if match:
220 return match.group(1).strip()
221 return ""
223 def _extract_code_block(self, text: str) -> str:
224 """Extract Solidity code block from markdown"""
225 pattern = r"```(?:solidity)?\n(.*?)```"
226 match = re.search(pattern, text, re.DOTALL)
227 if match:
228 return match.group(1).strip()
229 return ""
232# Convenience function for easy integration
233def enhance_if_available(vuln: Dict, code: str) -> Any:
234 """
235 Try to enhance with LLM, fall back gracefully if unavailable
237 Returns original vuln dict if LLM not available, EnhancedVulnerability otherwise
238 """
239 try:
240 enhancer = LLMEnhancer()
241 enhanced = enhancer.enhance_vulnerability(vuln, code)
242 return enhanced.to_dict()
243 except (ValueError, ImportError, Exception) as e:
244 logger.debug(f"LLM enhancement unavailable: {e}")
245 return vuln