Coverage for src/alprina_cli/agents/web3_auditor/defi_risk_assessor.py: 16%
170 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-14 11:27 +0100
1"""
2DeFi Economic Risk Assessor with AI Enhancement
4Analyzes smart contracts for economic attack vectors that traditional static analysis misses.
5Uses both pattern matching and LLM-powered contextual analysis.
6"""
8import re
9from typing import List, Dict, Any, Optional
10from dataclasses import dataclass
11from enum import Enum
13class EconomicRiskType(Enum):
14 FLASH_LOAN_ATTACK = "flash_loan_attack"
15 PRICE_ORACLE_MANIPULATION = "price_oracle_manipulation"
16 LIQUIDITY_DRAIN = "liquidity_drain"
17 MEV_EXTRACTION = "mev_extraction"
18 YIELD_FARMING_EXPLOIT = "yield_farming_exploit"
19 CROSS_CHAIN_BRIDGE = "cross_chain_bridge"
20 TOKEN_MINT_MANIPULATION = "token_mint_manipulation"
21 GOVERNANCE_ATTACK = "governance_attack"
23@dataclass
24class EconomicRisk:
25 """Represents an economic risk in a DeFi protocol"""
26 risk_type: EconomicRiskType
27 severity: str # "critical", "high", "medium", "low"
28 title: str
29 description: str
30 attack_scenario: str
31 file_path: str
32 line_number: Optional[int]
33 contract_name: str
34 economic_impact: str
35 remediation: str
36 confidence: int = 80 # 0-100
38class DeFiRiskAssessor:
39 """
40 AI-powered DeFi economic risk assessor
41 Goes beyond code vulnerabilities to detect business logic flaws
42 """
44 def __init__(self):
45 self.risk_patterns = self._initialize_risk_patterns()
46 self.llm_client = self._initialize_llm_client()
48 def assess_economic_risks(self, contract_code: str, protocol_context: Dict[str, Any]) -> List[EconomicRisk]:
49 """
50 Comprehensive economic risk assessment for DeFi protocols
52 Args:
53 contract_code: Solidity or other blockchain contract code
54 protocol_context: Information about the DeFi protocol
56 Returns:
57 List of detected economic risks
58 """
59 risks = []
61 # Traditional pattern-based analysis
62 flash_loan_risks = self._detect_flash_loan_vulnerabilities(contract_code, protocol_context)
63 oracle_risks = self._detect_price_oracle_risks(contract_code, protocol_context)
64 liquidity_risks = self._detect_liquidity_drain_risks(contract_code, protocol_context)
65 governance_risks = self._detect_governance_attacks(contract_code, protocol_context)
67 risks.extend(flash_loan_risks)
68 risks.extend(oracle_risks)
69 risks.extend(liquidity_risks)
70 risks.extend(governance_risks)
72 # AI-enhanced analysis for complex patterns
73 llm_risks = self._llm_economic_analysis(contract_code, protocol_context)
74 risks.extend(llm_risks)
76 return risks
78 def _detect_flash_loan_vulnerabilities(self, contract_code: str, context: Dict[str, Any]) -> List[EconomicRisk]:
79 """Detect flash loan attack vectors"""
80 risks = []
81 lines = contract_code.split('\n')
83 for i, line in enumerate(lines):
84 line_content = line.strip()
86 # Pattern 1: Functions that update state based on token balances
87 if any(pattern in line_content for pattern in ['balanceOf', 'totalSupply', 'reserve']):
88 # Check if this is a liquidity function
89 next_lines = lines[i+1:i+10] # Look at next several lines
90 function_block = '\n'.join(next_lines)
92 if 'update' in function_block.lower() or 'calculate' in function_block.lower():
93 # Look for missing reentrancy checks or price validation
94 has_protection = any(
95 protection in function_block.lower()
96 for protection in ['reentrancyguard', 'nonreentrant', 'require(msg.value']
97 )
99 if not has_protection:
100 risk = EconomicRisk(
101 risk_type=EconomicRiskType.FLASH_LOAN_ATTACK,
102 severity="high",
103 title="Flash Loan Manipulation Risk",
104 description="Function updates state based on balances without flash loan protection",
105 attack_scenario="Attacker takes flash loan → Manipulates token prices → Exploits price-dependent calculation → Repays loan with profit",
106 file_path=context.get('file_path', 'unknown'),
107 line_number=i + 1,
108 contract_name=context.get('contract_name', 'unknown'),
109 economic_impact="Potential protocol fund drain if exploited",
110 remediation="Add reentrancy protection and validate price calculations with time-weighted averages",
111 confidence=75
112 )
113 risks.append(risk)
115 # Pattern 2: Liquidity provision functions with price calculation
116 liquidity_patterns = ['addLiquidity', 'removeLiquidity', 'swap', 'exchange']
117 for pattern in liquidity_patterns:
118 context_window = contract_code[max(0, i-5):i+5]
119 if pattern in context_window and 'calculate' in context_window:
120 # Check for price oracle usage
121 if 'uniswap' in context_window.lower() or 'getAmountsOut' in context_window:
122 risk = EconomicRisk(
123 risk_type=EconomicRiskType.FLASH_LOAN_ATTACK,
124 severity="high",
125 title="Flash Loan Oracle Manipulation",
126 description="Liquidity function uses on-chain price oracle vulnerable to manipulation",
127 attack_scenario="Attaker uses flash loan to manipulate exchange rate → Drains liquidity pool → Repays loan",
128 file_path=context.get('file_path', 'unknown'),
129 line_number=i + 1,
130 contract_name=context.get('contract_name', 'unknown'),
131 economic_impact="Complete liquidity pool drain possible",
132 remediation="Use TWAP price oracle or validate against multiple sources",
133 confidence=80
134 )
135 risks.append(risk)
137 return risks
139 def _detect_price_oracle_risks(self, contract_code: str, context: Dict[str, Any]) -> List[EconomicRisk]:
140 """Detect price oracle manipulation vulnerabilities"""
141 risks = []
142 lines = contract_code.split('\n')
144 oracle_sources = ['uniswap', 'sushiswap', 'curve', 'balancer', 'getAmountsOut']
146 for i, line in enumerate(lines):
147 line_content = line.strip()
149 # Check for direct oracle usage without validation
150 for oracle in oracle_sources:
151 if oracle in line_content.lower():
152 # Look at surrounding code for price usage
153 context_lines = lines[max(0, i-3):i+8] # 3 lines before, 7 after
154 context_block = '\n'.join(context_lines)
156 # Check for price-dependent calculations
157 if any(calc in context_block for calc in ['amount', 'price', 'value', 'calculate']):
158 # Check for validation
159 has_validation = any(
160 validation in context_block.lower()
161 for validation in ['min', 'max', '>=', '<=', 'require']
162 )
164 if not has_validation:
165 risk = EconomicRisk(
166 risk_type=EconomicRiskType.PRICE_ORACLE_MANIPULATION,
167 severity="high",
168 title="Unvalidated Price Oracle",
169 description=f"Using {oracle.title()} price oracle without validation",
170 attack_scenario="Attacker manipulates pool reserves → Invalid price data → Protocol miscalculates → Economic loss",
171 file_path=context.get('file_path', 'unknown'),
172 line_number=i + 1,
173 contract_name=context.get('contract_name', 'unknown'),
174 economic_impact="Protocol calculates wrong prices, leading to fund loss",
175 remediation="Implement price deviation checks and use multiple oracle sources",
176 confidence=85
177 )
178 risks.append(risk)
180 return risks
182 def _detect_liquidity_drain_risks(self, contract_code: str, context: Dict[str, Any]) -> List[EconomicRisk]:
183 """Detect liquidity draining attack vectors"""
184 risks = []
185 lines = contract_code.split('\n')
187 # Pattern: Functions that provide liquidity without restrictions
188 liquidity_patterns = [
189 'function.*withdraw',
190 'function.*redeem',
191 'function.*claim',
192 'function.*removeLiquidity'
193 ]
195 for i, line in enumerate(lines):
196 line_content = line.strip()
198 for pattern in liquidity_patterns:
199 if re.search(pattern, line_content):
200 # Look for access controls
201 function_context = lines[max(0, i-2):i+20] # Function scope
202 function_block = '\n'.join(function_context)
204 has_auth = any(
205 auth in function_block.lower()
206 for auth in ['onlyowner', 'require', 'modifiers']
207 )
209 has_amount_validation = any(
210 validation in function_block.lower()
211 for validation in ['require.*amount', 'min.*amount', '< balance']
212 )
214 if not has_auth or not has_amount_validation:
215 risk = EconomicRisk(
216 risk_type=EconomicRiskType.LIQUIDITY_DRAIN,
217 severity="critical",
218 title="Liquidity Drain Vulnerability",
219 description="Liquidity withdrawal function lacking proper controls",
220 attack_scenario="Attacker calls unauthorized withdrawal → Drains all available liquidity → Protocol becomes insolvent",
221 file_path=context.get('file_path', 'unknown'),
222 line_number=i + 1,
223 contract_name=context.get('contract_name', 'unknown'),
224 economic_impact="Complete fund loss through liquidity drain",
225 remediation="Add proper access controls, amount validations, and withdrawal limits",
226 confidence=90
227 )
228 risks.append(risk)
230 return risks
232 def _detect_governance_attacks(self, contract_code: str, context: Dict[str, Any]) -> List[EconomicRisk]:
233 """Detect governance system attack vectors"""
234 risks = []
235 lines = contract_code.split('\n')
237 governance_patterns = [
238 'function vote',
239 'function propose',
240 'function execute',
241 'function delegate'
242 ]
244 for i, line in enumerate(lines):
245 line_content = line.strip()
247 for pattern in governance_patterns:
248 if pattern in line_content:
249 # Check for voting power manipulation risks
250 context_lines = lines[max(0, i-3):i+15] # Goverance function scope
251 context_block = '\n'.join(context_lines)
253 # Look for token balance usage without time-lock
254 if ('balance' in context_block and 'vote' in context_block.lower()):
255 has_timelock = any(
256 timelock in context_block.lower()
257 for timelock in ['timelock', 'delay', 'waiting', 'period']
258 )
260 if not has_timelock:
261 risk = EconomicRisk(
262 risk_type=EconomicRiskType.GOVERNANCE_ATTACK,
263 severity="high",
264 title="Governance Flash Loan Attack",
265 description="Voting mechanism vulnerable to flash loan manipulation",
266 attack_scenario="Attaker takes flash loan → Gains voting power → Passes malicious proposal → Executes against protocol → Repays loan",
267 file_path=context.get('file_path', 'unknown'),
268 line_number=i + 1,
269 contract_name=context.get('contract_name', 'unknown'),
270 economic_impact="Protocol governance can be seized, leading to full protocol control",
271 remediation="Implement voting power time-locks and minimum holding periods",
272 confidence=80
273 )
274 risks.append(risk)
276 return risks
278 def _llm_economic_analysis(self, contract_code: str, context: Dict[str, Any]) -> List[EconomicRisk]:
279 """AI-powered economic risk analysis for complex attack vectors"""
280 risks = []
282 try:
283 # Prepare context for LLM analysis
284 protocol_type = context.get('protocol_type', 'DeFi Protocol')
285 contract_name = context.get('contract_name', 'Smart Contract')
287 prompt = f"""
288 Analyze this DeFi smart contract for complex economic attack vectors that traditional static analysis might miss:
290 Protocol Type: {protocol_type}
291 Contract Name: {contract_name}
293 Contract Code:
294 {contract_code}
296 Focus on these economic attack patterns:
297 1. Flash Loan Attack Vectors - identify where price calculations could be manipulated
298 2. Cross-Protocol Arbitrage Exploits - look for interactions with other DeFi protocols
299 3. MEV (Maximal Extractable Value) opportunities - identify transaction ordering exploits
300 4. Token Supply/Governance Manipulation - check for voting/influence manipulation
301 5. Liquidity Manipulation Scenarios - identify pool manipulation possibilities
302 6. Cross-Chain Bridge Risks - if applicable, check for multi-chain vulnerabilities
304 For each risk identified, provide:
305 - Attack scenario in step-by-step format
306 - Economic impact assessment
307 - Mitigation strategies
309 Analyze deeply and think about complex multi-step attacks.
310 """
312 # Simulate LLM response (in production, this would call actual LLM)
313 llm_response = self._simulate_llm_analysis(contract_code, context)
315 if llm_response:
316 # Parse LLM response into EconomicRisk objects
317 llm_risks = self._parse_llm_response(llm_response, contract_code, context)
318 risks.extend(llm_risks)
320 except Exception as e:
321 # Add fallback risk if LLM analysis fails
322 fallback_risk = EconomicRisk(
323 risk_type=EconomicRiskType.LIQUIDITY_DRAIN,
324 severity="medium",
325 title="AI Analysis Limitation",
326 description=f"Complete economic analysis unavailable: {str(e)[:100]}",
327 attack_scenario="Consider manual security audit for complex economic vectors",
328 file_path=context.get('file_path', 'unknown'),
329 line_number=None,
330 contract_name=context.get('contract_name', 'unknown'),
331 economic_impact="Unknown - manual review recommended",
332 remediation="Manual security audit recommended for comprehensive economic analysis",
333 confidence=30
334 )
335 risks.append(fallback_risk)
337 return risks
339 def _simulate_llm_analysis(self, contract_code: str, context: Dict[str, Any]) -> Optional[str]:
340 """Simulate LLM analysis - in production, this would call actual LLM API"""
342 # Check for common DeFi patterns and return simulated responses
343 contract_lower = contract_code.lower()
345 if 'uniswap' in contract_lower and 'getamountsout' in contract_lower:
346 return """
347 RISK: Flash Loan Oracle Manipulation
348 Severity: High
350 Description: This contract uses Uniswap price oracle (getAmountsOut) without validation, making it vulnerable to flash loan manipulation attacks.
352 Attack Scenario:
353 1. Attaker identifies large reserves in Uniswap pool
354 2. Takes flash loan from Aave/Compound
355 3. Uses flash loan to manipulate pool prices
356 4. Calls protocol function using manipulated oracle prices
357 5. Protocol calculates incorrect prices, allowing arbitrage
358 6. Attaker repays flash loan and keeps profit
360 Economic Impact: Complete liquidity drain possible
362 Mitigation: Implement TWAP price oracle, validate maximum price deviations
363 """
365 elif 'token' in contract_lower and 'mint' in contract_lower:
366 return """
367 RISK: Token Supply Manipulation
368 Severity: High
370 Description: Minting function lacks proper access controls, potentially allowing unlimited token creation.
372 Attack Scenario:
373 1. Attaker identifies minting privilege escalation
374 2. Mints large number of tokens
375 3. Uses tokens to manipulate voting/proportions
376 4. Drains platform resources through voting power
377 5. Results in protocol fund loss
379 Economic Impact: Protocol governance and fund control loss
381 Mitigation: Add strict access controls, implement minting caps, add governance delays
382 """
384 elif 'withdraw' in contract_lower and 'balance' in contract_lower:
385 return """
386 RISK: Liquidity Drain Attack
387 Severity: Critical
389 Description: Withdrawal function lacks adequate access controls and balance validations.
391 Attack Scenario:
392 1. Attaker identifies unprotected withdrawal function
393 2. Calls withdrawal to drain available funds
394 3. Protocol becomes insolvent or drained completely
395 4. All other users lose access to funds
397 Economic Impact: Complete fund loss possible
399 Mitigation: Implement withdrawal limits, access controls, and balance validations
400 """
402 return None
404 def _parse_llm_response(self, llm_response: str, contract_code: str, context: Dict[str, Any]) -> List[EconomicRisk]:
405 """Parse LLM response into EconomicRisk objects"""
406 risks = []
408 try:
409 # Parse simulated response format
410 if 'RISK:' in llm_response:
411 sections = llm_response.split('RISK:')[1:]
413 for section in sections:
414 lines = section.strip().split('\n')
415 if len(lines) < 5:
416 continue
418 # Extract key information
419 title = lines[0].split(':')[1].strip() if ':' in lines[0] else "Economic Risk"
420 severity = lines[1].split(':')[1].strip().lower() if ':' in lines[1] else "medium"
421 description = lines[2].split(':')[1].strip() if ':' in lines[2] else "Economic vulnerability detected"
423 # Find attack scenario
424 attack_scenario = ""
425 description_start = False
426 for line in lines[3:]:
427 if line.strip().startswith("Attack Scenario:"):
428 description_start = True
429 attack_scenario = line.split("Attack Scenario:")[1].strip()
430 elif description_start:
431 if line.strip().startswith("Economic Impact:"):
432 break
433 attack_scenario += " " + line.strip()
435 # Determine risk type based on content
436 risk_type = EconomicRiskType.LIQUIDITY_DRAIN # default
437 if "flash loan" in llm_response.lower():
438 risk_type = EconomicRiskType.FLASH_LOAN_ATTACK
439 elif "oracle" in llm_response.lower() or "price" in llm_response.lower():
440 risk_type = EconomicRiskType.PRICE_ORACLE_MANIPULATION
441 elif "token" in llm_response.lower() or "mint" in llm_response.lower():
442 risk_type = EconomicRiskType.TOKEN_MINT_MANIPULATION
444 # Extract economic impact and mitigation
445 economic_impact = "Potential economic loss"
446 remediation = "Security audit recommended"
448 for line in lines:
449 if line.strip().startswith("Economic Impact:"):
450 economic_impact = line.split("Economic Impact:")[1].strip()
451 elif line.strip().startswith("Mitigation:"):
452 remediation = line.split("Mitigation:")[1].strip()
454 risk = EconomicRisk(
455 risk_type=risk_type,
456 severity=severity if severity in ["critical", "high", "medium", "low"] else "medium",
457 title=title,
458 description=description,
459 attack_scenario=attack_scenario.strip(),
460 file_path=context.get('file_path', 'unknown'),
461 line_number=None,
462 contract_name=context.get('contract_name', 'unknown'),
463 economic_impact=economic_impact,
464 remediation=remediation,
465 confidence=75 # LLM-based confidence
466 )
467 risks.append(risk)
469 except Exception as e:
470 # Debug: LLM parsing failed
471 pass
473 return risks
475 def _initialize_risk_patterns(self) -> Dict[str, List[str]]:
476 """Initialize economic risk pattern detectors"""
477 return {
478 'flash_loan': [
479 r'flashloan',
480 r'borrow.*flash',
481 r'uniswap.*price',
482 r'getAmountsOut'
483 ],
484 'oracle_manipulation': [
485 r'oracle',
486 r'price.*feed',
487 r'uniswap.*router',
488 r'chainlink'
489 ],
490 'liquidity_drain': [
491 r'withdraw.*all',
492 r'balance.*transfer',
493 r'drain.*liquidity'
494 ],
495 'governance': [
496 r'vote.*balance',
497 r'governance.*token',
498 r'propose.*execute',
499 r'timelock'
500 ]
501 }
503 def _initialize_llm_client(self):
504 """Initialize LLM client - in production, this would connect to Claude/ChatGPT"""
505 # For now, return None - LLM logic uses simulated responses
506 return None