Coverage for src/alprina_cli/agents/web3_auditor/web3_auditor.py: 17%
293 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-14 11:27 +0100
1"""
2Web3/DeFi Security Auditor Agent
4Enterprise-grade smart contract security analysis combining:
51. Static vulnerability detection (OWASP Smart Contract Top 10)
62. Economic risk assessment (AI-enhanced)
73. Multi-chain support (Ethereum, Solana, Polygon, BSC)
84. Flash loan attack vector analysis
95. Real-time DeFi protocol monitoring
10"""
12import asyncio
13import json
14from pathlib import Path
15from typing import Dict, Any, List, Optional, Tuple
16from dataclasses import dataclass, asdict
17from enum import Enum
19from .solidity_analyzer import SolidityStaticAnalyzer, SolidityVulnerability
20from .defi_risk_assessor import DeFiRiskAssessor, EconomicRisk, EconomicRiskType
21from .multi_chain_scanner import MultiChainScanner, BlockchainType
23# LLM Enhancement (optional)
24try:
25 from ..llm_enhancer import LLMEnhancer
26 from ..llm_config import LLMConfig
27 LLM_AVAILABLE = True
28except ImportError:
29 LLM_AVAILABLE = False
31class VulnerabilitySeverity(Enum):
32 CRITICAL = "critical"
33 HIGH = "high"
34 MEDIUM = "medium"
35 LOW = "low"
36 INFO = "info"
38@dataclass
39class Web3SecurityReport:
40 """Comprehensive Web3 security report"""
41 agent: str
42 status: str
43 blockchain_type: str
44 contracts_scanned: List[str]
45 vulnerabilities: List[Dict[str, Any]]
46 economic_risks: List[Dict[str, Any]]
47 risk_score: int
48 exploit_simulation: List[Dict[str, Any]]
49 recommendations: List[str]
50 summary: Dict[str, int]
51 confidence_score: float
53class Web3AuditorAgent:
54 """
55 Web3/DeFi Security Auditor - Startup-focused blockchain security analysis
57 Comprehensive security testing for smart contracts and DeFi protocols
58 """
60 def __init__(self):
61 self.name = "Web3/DeFi Security Auditor"
62 self.agent_type = "web3-security"
63 self.description = "Enterprise-grade Web3 smart contract security analysis with AI-enhanced economic risk assessment"
65 # Core analysis engines
66 self.solidity_analyzer = SolidityStaticAnalyzer()
67 self.defi_assessor = DeFiRiskAssessor()
68 self.multi_chain_scanner = MultiChainScanner()
70 # Vulnerability databases
71 self.exploit_database = self._initialize_exploit_database()
72 self.owasp_patterns = self._load_owasp_smart_contract_patterns()
74 # Supported blockchain platforms
75 self.supported_chains = [
76 BlockchainType.ETHEREUM,
77 BlockchainType.POLYGON,
78 BlockchainType.BSC,
79 BlockchainType.AVALANCHE,
80 BlockchainType.ARBITRUM,
81 BlockchainType.OPTIMISM
82 ]
84 def analyze_directory(self, directory_path: str) -> Web3SecurityReport:
85 """
86 Analyze a directory of smart contracts
88 Args:
89 directory_path: Path to directory containing smart contracts
91 Returns:
92 Comprehensive Web3 security report
93 """
94 directory = Path(directory_path)
96 if not directory.exists():
97 return self._error_report(f"Directory not found: {directory_path}")
99 # Discover smart contracts
100 contract_files = self._discover_contract_files(directory)
102 if not contract_files:
103 return self._error_report(f"No smart contracts found in {directory_path}")
105 all_vulnerabilities = []
106 all_economic_risks = []
107 contracts_scanned = []
109 # Analyze each contract file
110 for contract_file in contract_files:
111 try:
112 contract_report = self.analyze_contract_file(contract_file)
114 all_vulnerabilities.extend(contract_report.get('vulnerabilities', []))
115 all_economic_risks.extend(contract_report.get('economic_risks', []))
117 if contract_report.get('contracts'):
118 contracts_scanned.extend(contract_report['contracts'])
120 except Exception as e:
121 error_vuln = {
122 'severity': 'low',
123 'type': 'Analysis Error',
124 'title': f'Failed to analyze {contract_file}',
125 'description': str(e),
126 'file_path': str(contract_file),
127 'confidence': 20
128 }
129 all_vulnerabilities.append(error_vuln)
131 # Calculate risk scores
132 overall_risk_score = self._calculate_overall_risk_score(all_vulnerabilities, all_economic_risks)
134 # Generate recommendations
135 recommendations = self._generate_recommendations(all_vulnerabilities, all_economic_risks)
137 # Mock exploit simulation results
138 exploit_simulation = self._simulate_exploit_scenarios(all_vulnerabilities, all_economic_risks)
140 return Web3SecurityReport(
141 agent=self.name,
142 status="success",
143 blockchain_type="multi-chain",
144 contracts_scanned=list(set(contracts_scanned)),
145 vulnerabilities=self._format_vulnerabilities(all_vulnerabilities),
146 economic_risks=self._format_economic_risks(all_economic_risks),
147 risk_score=overall_risk_score,
148 exploit_simulation=exploit_simulation,
149 recommendations=recommendations,
150 summary= self._generate_summary(all_vulnerabilities, all_economic_risks),
151 confidence_score=self._calculate_confidence_score(all_vulnerabilities, all_economic_risks)
152 )
154 def analyze_contract_file(self, file_path: str, blockchain_type: str = "ethereum") -> Dict[str, Any]:
155 """
156 Analyze a single smart contract file
158 Args:
159 file_path: Path to contract file
160 blockchain_type: Target blockchain platform
162 Returns:
163 Analysis results for the contract
164 """
165 file_path = Path(file_path)
167 if not file_path.exists():
168 return {
169 'status': 'error',
170 'error': f'File not found: {file_path}',
171 'vulnerabilities': [],
172 'economic_risks': [],
173 'contracts': []
174 }
176 try:
177 contract_code = file_path.read_text(encoding='utf-8')
179 # Detect contract language
180 contract_language = self._detect_contract_language(file_path, contract_code)
182 if contract_language == "solidity":
183 return self._analyze_solidity_contract(contract_code, str(file_path))
184 elif contract_language == "solana":
185 return self._analyze_solana_contract(contract_code, str(file_path))
186 else:
187 return self._analyze_generic_contract(contract_code, str(file_path))
189 except Exception as e:
190 return {
191 'status': 'error',
192 'error': f'Failed to analyze {file_path}: {str(e)}',
193 'vulnerabilities': [],
194 'economic_risks': [],
195 'contracts': []
196 }
198 def _analyze_solidity_contract(self, contract_code: str, file_path: str) -> Dict[str, Any]:
199 """Analyze Solidity smart contract"""
201 # Static analysis
202 static_vulnerabilities = self.solidity_analyzer.analyze_contract(contract_code, file_path)
204 # Economic risk assessment
205 protocol_context = {
206 'file_path': file_path,
207 'language': 'solidity',
208 'protocol_type': self._infer_protocol_type(contract_code)
209 }
210 economic_risks = self.defi_assessor.assess_economic_risks(contract_code, protocol_context)
212 # Extract contract names
213 contract_names = self._extract_contract_names(contract_code)
215 return {
216 'status': 'success',
217 'language': 'solidity',
218 'contracts': contract_names,
219 'vulnerabilities': static_vulnerabilities,
220 'economic_risks': economic_risks
221 }
223 def _analyze_solana_contract(self, contract_code: str, file_path: str) -> Dict[str, Any]:
224 """Analyze Solana program (Rust)"""
226 # Placeholder for Solana analysis - would implement Rust static analysis
227 vulnerabilities = []
228 economic_risks = []
230 # Basic Solana-specific checks
231 rust_patterns = [
232 ('unsafe', 'code', 'high', 'Unsafe block detected - potential security risk'),
233 ('vec', 'memory', 'medium', 'Vector usage may cause memory issues'),
234 ('unwrap', 'error', 'medium', 'unwrap() calls may panic on error'),
235 ]
237 lines = contract_code.split('\n')
238 for i, line in enumerate(lines):
239 for pattern, vuln_type, severity, description in rust_patterns:
240 if pattern in line:
241 vulnerabilities.append({
242 'vulnerability_type': vuln_type,
243 'severity': severity,
244 'title': f'{vuln_type.title()} Risk',
245 'description': description,
246 'file_path': file_path,
247 'line_number': i + 1,
248 'contract_name': 'solana_program',
249 'code_snippet': line.strip(),
250 'confidence': 70
251 })
253 program_names = self._extract_rust_program_names(contract_code)
255 return {
256 'status': 'success',
257 'language': 'rust',
258 'contracts': program_names,
259 'vulnerabilities': vulnerabilities,
260 'economic_risks': economic_risks
261 }
263 def _analyze_generic_contract(self, contract_code: str, file_path: str) -> Dict[str, Any]:
264 """Analyze generic smart contract (fallback logic)"""
266 vulnerabilities = []
267 economic_risks = []
269 # Basic pattern matching for common contract patterns
270 suspicious_patterns = [
271 ('transfer', 'security', 'medium', 'Transfer function detected - manual audit recommended'),
272 ('balance', 'security', 'medium', 'Balance manipulation detected'),
273 ('owner', 'access', 'medium', 'Owner function - verify access controls'),
274 ]
276 lines = contract_code.split('\n')
277 for i, line in enumerate(lines):
278 for pattern, category, severity, description in suspicious_patterns:
279 if pattern in line.lower():
280 vulnerabilities.append({
281 'vulnerability_type': category,
282 'severity': severity,
283 'title': f'{category.title()} Pattern',
284 'description': description,
285 'file_path': file_path,
286 'line_number': i + 1,
287 'contract_name': 'unknown',
288 'code_snippet': line.strip(),
289 'confidence': 50
290 })
292 return {
293 'status': 'success',
294 'language': 'unknown',
295 'contracts': ['generic_contract'],
296 'vulnerabilities': vulnerabilities,
297 'economic_risks': economic_risks
298 }
300 def _discover_contract_files(self, directory: Path) -> List[Path]:
301 """Discover smart contract files in directory"""
302 contract_files = []
304 # Solidity files
305 sol_patterns = list(directory.rglob("*.sol"))
306 contract_files.extend(sol_patterns)
308 # Solana/Rust files
309 rs_patterns = []
310 for rust_file in directory.rglob("*.rs"):
311 # Check if it's likely a Solana program
312 rust_content = rust_file.read_text(encoding='utf-8', errors='ignore')
313 if any(solana_keyword in rust_content for solana_keyword in ['sol_program', 'account_info', 'pubkey']):
314 rs_patterns.append(rust_file)
315 contract_files.extend(rs_patterns)
317 # Move files
318 move_patterns = list(directory.rglob("*.move"))
319 contract_files.extend(move_patterns)
321 return contract_files
323 def _detect_contract_language(self, file_path: Path, contract_code: str) -> str:
324 """Detect smart contract programming language"""
326 file_extension = file_path.suffix.lower()
328 if file_extension == '.sol':
329 return 'solidity'
330 elif file_extension == '.rs':
331 if 'sol_program' in contract_code or 'account_info' in contract_code:
332 return 'solana'
333 else:
334 return 'rust'
335 elif file_extension == '.move':
336 return 'move'
337 else:
338 # Try to detect from content
339 if any(keyword in contract_code for keyword in ['contract', 'function', 'modifier']):
340 return 'solidity'
341 elif any(keyword in contract_code for keyword in ['program', 'account_info', 'pubkey']):
342 return 'solana'
343 else:
344 return 'unknown'
346 def _extract_contract_names(self, contract_code: str) -> List[str]:
347 """Extract contract names from Solidity code"""
348 import re
349 contract_pattern = r'(abstract\s+)?contract\s+(\w+)'
350 contracts = re.findall(contract_pattern, contract_code)
351 return [name for _, name in contracts]
353 def _extract_rust_program_names(self, contract_code: str) -> List[str]:
354 """Extract program names from Solana Rust code"""
355 import re
356 program_pattern = r'sol_program!\s*\[(.*?)\]\s*fn\s*(\w+)'
357 programs = re.findall(program_pattern, contract_code)
358 return [name for _, name in programs]
360 def _infer_protocol_type(self, contract_code: str) -> str:
361 """Infer DeFi protocol type from contract code"""
362 contract_lower = contract_code.lower()
364 if any(token in contract_lower for token in ['uniswap', 'swap', 'liquidity']):
365 return 'Decentralized Exchange (DEX)'
366 elif any(token in contract_lower for token in ['lending', 'borrow', 'interest']):
367 return 'Lending Protocol'
368 elif any(token in contract_lower for token in ['price', 'oracle', 'feed']):
369 return 'Price Oracle'
370 elif any(token in contract_lower for token in ['governance', 'vote', 'propose']):
371 return 'Governance Protocol'
372 elif any(token in contract_lower for token in ['token', 'mint', 'burn']):
373 return 'Token Contract'
374 else:
375 return 'General DeFi Protocol'
377 def _calculate_overall_risk_score(self, vulnerabilities: List, economic_risks: List) -> int:
378 """Calculate overall risk score 0-100"""
379 if not vulnerabilities and not economic_risks:
380 return 0
382 severity_weights = {
383 'critical': 25,
384 'high': 15,
385 'medium': 8,
386 'low': 3
387 }
389 # Calculate vulnerability score
390 vulnerability_score = sum(
391 severity_weights.get(vuln.get('severity', 'low'), 3)
392 for vuln in vulnerabilities
393 )
395 # Economic risks are weighted higher as they typically have bigger impact
396 economic_risk_score = sum(
397 severity_weights.get(risk.get('severity', 'medium'), 8) * 1.2
398 for risk in economic_risks
399 )
401 total_score = vulnerability_score + economic_risk_score
402 return min(100, int(total_score))
404 def _generate_recommendations(self, vulnerabilities: List, economic_risks: List) -> List[str]:
405 """Generate actionable security recommendations"""
406 recommendations = []
408 # Based on vulnerability types
409 vuln_types = {v.get('vulnerability_type') for v in vulnerabilities}
411 if 'reentrancy' in vuln_types:
412 recommendations.append("Implement ReentrancyGuard modifier and follow checks-effects-interactions pattern")
414 if 'access_control' in vuln_types:
415 recommendations.append("Review and implement proper access control for all critical functions")
417 if 'integer_overflow_underflow' in vuln_types:
418 recommendations.append("Use SafeMath library or Solidity 0.8+ which has built-in overflow protection")
420 # Based on economic risk types
421 economic_types = {r.get('risk_type') for r in economic_risks}
423 if EconomicRiskType.FLASH_LOAN_ATTACK.value in economic_types:
424 recommendations.append("Implement flash loan protection and validate price calculations")
426 if EconomicRiskType.PRICE_ORACLE_MANIPULATION.value in economic_types:
427 recommendations.append("Use multiple oracle sources and implement price deviation checks")
429 if EconomicRiskType.LIQUIDITY_DRAIN.value in economic_types:
430 recommendations.append("Add withdrawal limits and implement emergency pause mechanisms")
432 # General recommendations
433 recommendations.extend([
434 "Consider professional security audit before mainnet deployment",
435 "Implement comprehensive testing including fuzzing and property-based testing",
436 "Set up monitoring and alerting for suspicious activities"
437 ])
439 return list(set(recommendations)) # Remove duplicates
441 def _simulate_exploit_scenarios(self, vulnerabilities: List, economic_risks: List) -> List[Dict[str, Any]]:
442 """Simulate potential exploit scenarios"""
443 scenarios = []
445 # Flash loan attack simulation
446 flash_loan_vulns = [v for v in vulnerabilities if v.get('vulnerability_type') == 'reentrancy']
447 if flash_loan_vulns:
448 scenarios.append({
449 'attack_type': 'Flash Loan Manipulation',
450 'probability': 'High',
451 'potential_loss': 'Complete fund drain',
452 'description': 'Attacker uses flash loan to manipulate prices and exploit reentrancy',
453 'prevention': 'Add reentrancy protection and price validation'
454 })
456 # Governance attack simulation
457 governance_vulns = [v for v in vulnerabilities if v.get('vulnerability_type') == 'access_control']
458 if governance_vulns:
459 scenarios.append({
460 'attack_type': 'Governance Takeover',
461 'probability': 'Medium',
462 'potential_loss': 'Protocol control loss',
463 'description': 'Attaker gains voting power and passes malicious proposals',
464 'prevention': 'Implement voting time-locks and access controls'
465 })
467 return scenarios
469 def _generate_summary(self, vulnerabilities: List, economic_risks: List) -> Dict[str, int]:
470 """Generate summary statistics"""
471 summary = {
472 'critical': 0,
473 'high': 0,
474 'medium': 0,
475 'low': 0
476 }
478 for vuln in vulnerabilities:
479 severity = vuln.get('severity', 'low')
480 summary[severity] = summary.get(severity, 0) + 1
482 for risk in economic_risks:
483 severity = risk.get('severity', 'medium')
484 summary[severity] = summary.get(severity, 0) + 1
486 return summary
488 def _calculate_confidence_score(self, vulnerabilities: List, economic_risks: List) -> float:
489 """Calculate confidence score for analysis accuracy"""
490 if not vulnerabilities and not economic_risks:
491 return 100.0 # High confidence if no issues found
493 total_findings = len(vulnerabilities) + len(economic_risks)
494 confidence_sum = sum(v.get('confidence', 70) for v in vulnerabilities) + \
495 sum(r.get('confidence', 80) for r in economic_risks)
497 return min(100.0, round(confidence_sum / total_findings, 1))
499 def _format_vulnerabilities(self, vulnerabilities: List) -> List[Dict[str, Any]]:
500 """Format vulnerabilities for JSON serialization"""
501 return [v if isinstance(v, dict) else asdict(v) for v in vulnerabilities]
503 def _format_economic_risks(self, economic_risks: List) -> List[Dict[str, Any]]:
504 """Format economic risks for JSON serialization"""
505 return [r if isinstance(r, dict) else asdict(r) for r in economic_risks]
507 def _error_report(self, error_message: str) -> Web3SecurityReport:
508 """Generate error report"""
509 return Web3SecurityReport(
510 agent=self.name,
511 status="error",
512 blockchain_type="unknown",
513 contracts_scanned=[],
514 vulnerabilities=[],
515 economic_risks=[],
516 risk_score=0,
517 exploit_simulation=[],
518 recommendations=["Verify file path and ensure smart contracts are present"],
519 summary={'error': 1},
520 confidence_score=0.0
521 )
523 def _initialize_exploit_database(self) -> Dict[str, Any]:
524 """Initialize historical exploit database"""
525 return {
526 'flash_loan_attacks': {
527 "bzx_exploit_2020": "Attacker used flash loan to manipulate BZX price oracles",
528 "harvest_finance_2020": "Flash loan manipulation exploited pricing mechanism",
529 "pancakebunny_2022": "Flash loan attack on PancakeBunny liquidity pool"
530 },
531 'access_control_breaches': {
532 "parsec_finance_2022": "Access control bug allowed unlimited minting",
533 "iron_finance_2021": "Access control vulnerability in token contract"
534 },
535 'oracle_manipulation': {
536 "cream_finance_2021": "Oracle manipulation led to financial losses"
537 }
538 }
540 def _load_owasp_smart_contract_patterns(self) -> Dict[str, List[str]]:
541 """Load OWASP Smart Contract Top 10 patterns"""
542 return {
543 "reentrancy": ["call", "send", "transfer"],
544 "access_control": ["onlyOwner", "require", "if"],
545 "arithmetic": ["+", "-", "*", "/"],
546 "unchecked_calls": [".call(", ".delegatecall("],
547 "logic_errors": ["require", "revert", "assert"],
548 "oracles": ["uniswap", "chainlink", "price"],
549 "gas": ["gas", "block.gaslimit"]
550 }
552# Integration wrapper for existing agent system
553class Web3AuditorAgentWrapper:
554 """Integration wrapper for Web3/DeFi Security Auditor Agent"""
556 def __init__(self):
557 self.name = "Web3/DeFi Security Auditor"
558 self.agent_type = "web3-security"
559 self.description = "Enterprise-grade Web3 smart contract security analysis with AI-enhanced economic risk assessment"
560 self.auditor = Web3AuditorAgent()
562 # LLM enhancer (optional)
563 self.llm_enhancer = None
564 self.llm_enabled = False
566 if LLM_AVAILABLE:
567 try:
568 from loguru import logger
569 self.llm_enhancer = LLMEnhancer()
570 self.llm_enabled = True
571 logger.info("✅ LLM enhancement enabled (Claude AI)")
572 except Exception as e:
573 from loguru import logger
574 logger.info(f"LLM enhancement disabled: {e}")
576 def analyze(self, target: str, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
577 """
578 Analyze target for Web3 security vulnerabilities
580 Args:
581 target: Path to directory or file to analyze
582 options: Additional analysis options
584 Returns:
585 Dict containing analysis results
586 """
587 try:
588 target_path = Path(target)
590 if target_path.is_file():
591 # Single file analysis
592 blockchain_type = options.get('blockchain', 'ethereum') if options else 'ethereum'
593 result = self.auditor.analyze_contract_file(target, blockchain_type)
595 # Convert single file result to full report format
596 contracts = result.get('contracts', ['unknown'])
597 vulnerabilities = [v if isinstance(v, dict) else asdict(v)
598 for v in result.get('vulnerabilities', [])]
599 economic_risks = [r if isinstance(r, dict) else asdict(r)
600 for r in result.get('economic_risks', [])]
602 # ENHANCE with LLM if enabled
603 if self.llm_enabled and vulnerabilities:
604 vulnerabilities = self._enhance_vulnerabilities_with_llm(
605 vulnerabilities,
606 target_path.read_text(encoding='utf-8')
607 )
609 # Count LLM-enhanced vulnerabilities
610 llm_enhanced_count = sum(
611 1 for v in vulnerabilities if v.get('llm_enhanced', False)
612 )
614 return {
615 'agent': self.name,
616 'status': result.get('status', 'error'),
617 'blockchain_type': blockchain_type,
618 'contracts_scanned': contracts,
619 'risk_score': self._calculate_single_file_risk_score(vulnerabilities, economic_risks),
620 'vulnerabilities_count': len(vulnerabilities),
621 'economic_risks_count': len(economic_risks),
622 'vulnerabilities': vulnerabilities,
623 'economic_risks': economic_risks,
624 'summary': self._generate_single_file_summary(vulnerabilities, economic_risks),
625 'recommendations': self._generate_recommendations(vulnerabilities, economic_risks),
626 'llm_enhanced': self.llm_enabled,
627 'llm_enhanced_count': llm_enhanced_count
628 }
630 else:
631 # Directory analysis
632 report = self.auditor.analyze_directory(target)
634 return {
635 'agent': self.name,
636 'status': report.status,
637 'blockchain_type': report.blockchain_type,
638 'contracts_scanned': report.contracts_scanned,
639 'risk_score': report.risk_score,
640 'vulnerabilities_count': len(report.vulnerabilities),
641 'economic_risks_count': len(report.economic_risks),
642 'vulnerabilities': report.vulnerabilities,
643 'economic_risks': report.economic_risks,
644 'summary': report.summary,
645 'recommendations': report.recommendations,
646 'confidence_score': report.confidence_score,
647 'exploit_simulation': report.exploit_simulation
648 }
650 except Exception as e:
651 return {
652 'agent': self.name,
653 'status': 'error',
654 'error': str(e),
655 'vulnerabilities_count': 0,
656 'economic_risks_count': 0,
657 'risk_score': 0
658 }
660 def _calculate_single_file_risk_score(self, vulnerabilities: List, economic_risks: List) -> int:
661 """Calculate risk score for single file analysis"""
662 if not vulnerabilities and not economic_risks:
663 return 0
665 weights = {'critical': 25, 'high': 15, 'medium': 8, 'low': 3}
666 vuln_score = sum(weights.get(v.get('severity', 'low'), 3) for v in vulnerabilities)
667 econ_score = sum(weights.get(r.get('severity', 'medium'), 8) * 1.2 for r in economic_risks)
669 return min(100, int(vuln_score + econ_score))
671 def _generate_single_file_summary(self, vulnerabilities: List, economic_risks: List) -> Dict[str, int]:
672 """Generate summary for single file analysis"""
673 summary = {'critical': 0, 'high': 0, 'medium': 0, 'low': 0}
675 for vuln in vulnerabilities:
676 severity = vuln.get('severity', 'low')
677 summary[severity] = summary.get(severity, 0) + 1
679 for risk in economic_risks:
680 severity = risk.get('severity', 'medium')
681 summary[severity] = summary.get(severity, 0) + 1
683 return summary
685 def _generate_recommendations(self, vulnerabilities: List, economic_risks: List) -> List[str]:
686 """Generate recommendations for single file analysis"""
687 recommendations = []
689 if any(v.get('vulnerability_type') == 'reentrancy' for v in vulnerabilities):
690 recommendations.append("Implement ReentrancyGuard modifier")
692 if any(v.get('vulnerability_type') == 'access_control' for v in vulnerabilities):
693 recommendations.append("Review and fix access control mechanisms")
695 if any(r.get('risk_type') in ['flash_loan_attack', 'price_oracle_manipulation'] for r in economic_risks):
696 recommendations.append("Implement flash loan protection and oracle validation")
698 recommendations.extend([
699 "Consider professional security audit",
700 "Add comprehensive testing",
701 "Implement monitoring and alerting"
702 ])
704 return list(set(recommendations))
706 def _enhance_vulnerabilities_with_llm(
707 self,
708 vulnerabilities: List[Dict[str, Any]],
709 contract_code: str
710 ) -> List[Dict[str, Any]]:
711 """
712 Enhance vulnerabilities with LLM analysis
714 Args:
715 vulnerabilities: List of vulnerability dictionaries
716 contract_code: Full contract source code
718 Returns:
719 Enhanced vulnerability list
720 """
721 from loguru import logger
723 if not self.llm_enhancer or not vulnerabilities:
724 return vulnerabilities
726 # Sort vulnerabilities by severity (critical/high first)
727 severity_order = {'critical': 0, 'high': 1, 'medium': 2, 'low': 3}
728 sorted_vulns = sorted(
729 vulnerabilities,
730 key=lambda v: severity_order.get(v.get('severity', 'low'), 4)
731 )
733 enhanced_vulns = []
734 enhanced_count = 0
736 for vuln in vulnerabilities:
737 # Only enhance top 5 critical/high vulnerabilities (cost optimization)
738 should_enhance = (
739 vuln in sorted_vulns[:LLMConfig.MAX_VULNS_TO_ENHANCE] and
740 vuln.get('severity', '').lower() in ['critical', 'high', 'medium']
741 )
743 if should_enhance and enhanced_count < LLMConfig.MAX_VULNS_TO_ENHANCE:
744 try:
745 # Enhance with LLM
746 enhanced = self.llm_enhancer.enhance_vulnerability(vuln, contract_code)
747 enhanced_dict = enhanced.to_dict()
748 enhanced_vulns.append(enhanced_dict)
749 enhanced_count += 1
750 logger.debug(f"✅ Enhanced: {vuln.get('title', 'Unknown')}")
751 except Exception as e:
752 logger.warning(f"LLM enhancement failed for {vuln.get('title', 'Unknown')}: {e}")
753 enhanced_vulns.append(vuln)
754 else:
755 # Keep original vulnerability without enhancement
756 enhanced_vulns.append(vuln)
758 logger.info(f"LLM enhanced {enhanced_count}/{len(vulnerabilities)} vulnerabilities")
759 return enhanced_vulns