security_manager.py

#!/usr/bin/env python3
"""
Security Manager Module - Transaction Security and Risk Assessment
Similar to Rabby's security features: transaction simulation, risk alerts, approval management
"""

import hashlib
import hmac
import time
import json
import re
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
from enum import Enum
import secrets
import base64

try:
    from cryptography.fernet import Fernet
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    CRYPTO_AVAILABLE = True
except ImportError:
    CRYPTO_AVAILABLE = False
    print("Warning: cryptography not available. Install with: pip install cryptography")

class RiskLevel(Enum):
    """Risk assessment levels"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"

class SecurityAlert(Enum):
    """Types of security alerts"""
    UNKNOWN_CONTRACT = "unknown_contract"
    HIGH_VALUE_TRANSACTION = "high_value_transaction"
    SUSPICIOUS_RECIPIENT = "suspicious_recipient"
    UNLIMITED_APPROVAL = "unlimited_approval"
    KNOWN_SCAM = "known_scam"
    PHISHING_SITE = "phishing_site"
    RISKY_FUNCTION = "risky_function"
    UNUSUAL_GAS = "unusual_gas"

@dataclass
class SecurityAssessment:
    """Security assessment result"""
    risk_level: RiskLevel
    alerts: List[SecurityAlert]
    warnings: List[str]
    recommendations: List[str]
    score: int  # 0-100, lower is safer
    details: Dict[str, Any]

@dataclass
class ContractInfo:
    """Smart contract information"""
    address: str
    name: Optional[str]
    symbol: Optional[str]
    is_verified: bool
    is_proxy: bool
    creation_block: Optional[int]
    creator_address: Optional[str]
    source_code_available: bool
    audit_reports: List[str]
    risk_flags: List[str]

@dataclass
class ApprovalInfo:
    """Token approval information"""
    token_address: str
    token_symbol: str
    spender_address: str
    spender_name: Optional[str]
    approved_amount: str
    is_unlimited: bool
    last_activity: float
    risk_level: RiskLevel

class SecurityManager:
    """Manages security features and risk assessment"""
    
    def __init__(self):
        self.known_contracts = self._load_known_contracts()
        self.scam_addresses = self._load_scam_addresses()
        self.phishing_domains = self._load_phishing_domains()
        self.risky_functions = self._load_risky_functions()
        self.approval_cache = {}
        self.whitelist = set()
        self.blacklist = set()
        
        # Encryption setup
        self.encryption_key = None
        self.salt = None
        
    def _load_known_contracts(self) -> Dict[str, Dict[str, Any]]:
        """Load database of known contracts"""
        # In production, this would load from a database or API
        return {
            "0xa0b86a33e6417b8e.toLowerCase()": {
                "name": "Uniswap V2 Router",
                "type": "dex",
                "risk_level": "low",
                "verified": True
            },
            "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
                "name": "Uniswap V2 Router",
                "type": "dex", 
                "risk_level": "low",
                "verified": True
            },
            "0x1f9840a85d5af5bf1d1762f925bdaddc4201f984": {
                "name": "Uniswap Token",
                "type": "token",
                "risk_level": "low",
                "verified": True
            }
        }
    
    def _load_scam_addresses(self) -> set:
        """Load database of known scam addresses"""
        # In production, this would be updated from security feeds
        return {
            "0x000000000000000000000000000000000000dead",
            "0x0000000000000000000000000000000000000000"
        }
    
    def _load_phishing_domains(self) -> set:
        """Load database of phishing domains"""
        return {
            "unisvvap.org",
            "metamask-security.com",
            "etherscan-security.io"
        }
    
    def _load_risky_functions(self) -> Dict[str, Dict[str, Any]]:
        """Load database of risky contract functions"""
        return {
            "0xa9059cbb": {  # transfer(address,uint256)
                "name": "transfer",
                "risk": "low",
                "description": "Standard token transfer"
            },
            "0x095ea7b3": {  # approve(address,uint256)
                "name": "approve",
                "risk": "medium",
                "description": "Token approval - check amount"
            },
            "0x23b872dd": {  # transferFrom(address,address,uint256)
                "name": "transferFrom",
                "risk": "medium",
                "description": "Transfer from another address"
            },
            "0x2e1a7d4d": {  # withdraw(uint256)
                "name": "withdraw",
                "risk": "high",
                "description": "Withdrawal function"
            }
        }
    
    def setup_encryption(self, password: str) -> bool:
        """Setup encryption for sensitive data"""
        if not CRYPTO_AVAILABLE:
            return False
        
        try:
            # Generate salt
            self.salt = secrets.token_bytes(16)
            
            # Derive key from password
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=self.salt,
                iterations=100000,
            )
            key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
            self.encryption_key = Fernet(key)
            
            return True
            
        except Exception as e:
            print(f"Failed to setup encryption: {e}")
            return False
    
    def encrypt_data(self, data: str) -> Optional[str]:
        """Encrypt sensitive data"""
        if not self.encryption_key:
            return None
        
        try:
            encrypted = self.encryption_key.encrypt(data.encode())
            return base64.urlsafe_b64encode(encrypted).decode()
        except Exception:
            return None
    
    def decrypt_data(self, encrypted_data: str) -> Optional[str]:
        """Decrypt sensitive data"""
        if not self.encryption_key:
            return None
        
        try:
            encrypted_bytes = base64.urlsafe_b64decode(encrypted_data.encode())
            decrypted = self.encryption_key.decrypt(encrypted_bytes)
            return decrypted.decode()
        except Exception:
            return None
    
    def assess_transaction_risk(self, transaction: Dict[str, Any], 
                               network_config) -> SecurityAssessment:
        """
        Comprehensive transaction risk assessment
        Similar to Rabby's security engine
        """
        alerts = []
        warnings = []
        recommendations = []
        risk_score = 0
        details = {}
        
        # Analyze recipient address
        to_address = transaction.get('to', '').lower()
        
        if to_address in self.scam_addresses:
            alerts.append(SecurityAlert.KNOWN_SCAM)
            risk_score += 50
            warnings.append("Recipient is a known scam address")
            recommendations.append("Do not proceed with this transaction")
        
        if to_address in self.blacklist:
            alerts.append(SecurityAlert.SUSPICIOUS_RECIPIENT)
            risk_score += 30
            warnings.append("Recipient is on your blacklist")
        
        # Analyze transaction value
        value = int(transaction.get('value', 0))
        if value > 0:
            value_eth = value / 10**18
            
            if value_eth > 10:  # High value threshold
                alerts.append(SecurityAlert.HIGH_VALUE_TRANSACTION)
                risk_score += 20
                warnings.append(f"High value transaction: {value_eth:.4f} ETH")
                recommendations.append("Double-check recipient address")
        
        # Analyze contract interaction
        data = transaction.get('data', '0x')
        if data and data != '0x':
            function_sig = data[:10] if len(data) >= 10 else None
            
            if function_sig in self.risky_functions:
                func_info = self.risky_functions[function_sig]
                
                if func_info['risk'] == 'high':
                    alerts.append(SecurityAlert.RISKY_FUNCTION)
                    risk_score += 25
                    warnings.append(f"Calling risky function: {func_info['name']}")
                
                details['function_called'] = func_info
        
        # Check if interacting with unknown contract
        if to_address and to_address not in self.known_contracts:
            contract_info = self._analyze_contract(to_address, network_config)
            
            if not contract_info.is_verified:
                alerts.append(SecurityAlert.UNKNOWN_CONTRACT)
                risk_score += 15
                warnings.append("Interacting with unverified contract")
                recommendations.append("Verify contract source code before proceeding")
        
        # Gas analysis
        gas_price = transaction.get('gasPrice', 0)
        gas_limit = transaction.get('gas', 0)
        
        if gas_price > 100 * 10**9:  # > 100 Gwei
            alerts.append(SecurityAlert.UNUSUAL_GAS)
            risk_score += 10
            warnings.append("Unusually high gas price")
        
        # Determine risk level
        if risk_score >= 70:
            risk_level = RiskLevel.CRITICAL
        elif risk_score >= 40:
            risk_level = RiskLevel.HIGH
        elif risk_score >= 20:
            risk_level = RiskLevel.MEDIUM
        else:
            risk_level = RiskLevel.LOW
        
        # Add general recommendations
        if risk_level in [RiskLevel.HIGH, RiskLevel.CRITICAL]:
            recommendations.append("Consider canceling this transaction")
            recommendations.append("Verify all transaction details carefully")
        
        return SecurityAssessment(
            risk_level=risk_level,
            alerts=alerts,
            warnings=warnings,
            recommendations=recommendations,
            score=risk_score,
            details=details
        )
    
    def _analyze_contract(self, address: str, network_config) -> ContractInfo:
        """Analyze smart contract"""
        # In production, this would query blockchain explorers and contract databases
        return ContractInfo(
            address=address,
            name=None,
            symbol=None,
            is_verified=False,
            is_proxy=False,
            creation_block=None,
            creator_address=None,
            source_code_available=False,
            audit_reports=[],
            risk_flags=["unverified"]
        )
    
    def check_approval_risk(self, token_address: str, spender_address: str, 
                           amount: str) -> SecurityAssessment:
        """Check risk of token approval"""
        alerts = []
        warnings = []
        recommendations = []
        risk_score = 0
        
        # Check for unlimited approval
        max_uint256 = "115792089237316195423570985008687907853269984665640564039457584007913129639935"
        
        if amount == max_uint256 or int(amount) > 10**30:
            alerts.append(SecurityAlert.UNLIMITED_APPROVAL)
            risk_score += 30
            warnings.append("Unlimited token approval detected")
            recommendations.append("Consider setting a specific amount instead")
        
        # Check spender reputation
        if spender_address.lower() in self.scam_addresses:
            alerts.append(SecurityAlert.KNOWN_SCAM)
            risk_score += 50
            warnings.append("Approving tokens for known scam contract")
            recommendations.append("Do not approve - this is a scam")
        
        # Check if spender is unknown
        if spender_address.lower() not in self.known_contracts:
            alerts.append(SecurityAlert.UNKNOWN_CONTRACT)
            risk_score += 15
            warnings.append("Approving tokens for unknown contract")
            recommendations.append("Verify the contract before approving")
        
        # Determine risk level
        if risk_score >= 50:
            risk_level = RiskLevel.CRITICAL
        elif risk_score >= 30:
            risk_level = RiskLevel.HIGH
        elif risk_score >= 15:
            risk_level = RiskLevel.MEDIUM
        else:
            risk_level = RiskLevel.LOW
        
        return SecurityAssessment(
            risk_level=risk_level,
            alerts=alerts,
            warnings=warnings,
            recommendations=recommendations,
            score=risk_score,
            details={
                "token_address": token_address,
                "spender_address": spender_address,
                "amount": amount
            }
        )
    
    def simulate_transaction_preview(self, transaction: Dict[str, Any], 
                                   network_config) -> Dict[str, Any]:
        """
        Simulate transaction to show balance changes
        Similar to Rabby's transaction preview
        """
        # This would integrate with transaction simulation services
        # For now, return a mock preview
        
        preview = {
            "success": True,
            "balance_changes": [],
            "gas_used": transaction.get('gas', 21000),
            "events": [],
            "warnings": []
        }
        
        # Analyze potential balance changes
        value = int(transaction.get('value', 0))
        if value > 0:
            preview["balance_changes"].append({
                "type": "native_token",
                "address": transaction['from'],
                "change": f"-{value / 10**18:.6f} ETH",
                "direction": "out"
            })
            
            preview["balance_changes"].append({
                "type": "native_token", 
                "address": transaction['to'],
                "change": f"+{value / 10**18:.6f} ETH",
                "direction": "in"
            })
        
        return preview
    
    def check_website_security(self, url: str) -> Dict[str, Any]:
        """Check if website is potentially phishing"""
        domain = self._extract_domain(url)
        
        is_phishing = domain in self.phishing_domains
        
        # Check for common phishing patterns
        phishing_patterns = [
            r'.*meta.*mask.*',
            r'.*uniswap.*\.(?!org$)',
            r'.*ethereum.*security.*',
            r'.*wallet.*connect.*'
        ]
        
        suspicious_patterns = []
        for pattern in phishing_patterns:
            if re.match(pattern, domain, re.IGNORECASE):
                suspicious_patterns.append(pattern)
        
        risk_level = RiskLevel.HIGH if is_phishing or suspicious_patterns else RiskLevel.LOW
        
        return {
            "is_safe": not (is_phishing or suspicious_patterns),
            "is_known_phishing": is_phishing,
            "suspicious_patterns": suspicious_patterns,
            "risk_level": risk_level.value,
            "recommendation": "Do not connect wallet" if risk_level == RiskLevel.HIGH else "Appears safe"
        }
    
    def _extract_domain(self, url: str) -> str:
        """Extract domain from URL"""
        try:
            if not url.startswith(('http://', 'https://')):
                url = 'https://' + url
            
            from urllib.parse import urlparse
            parsed = urlparse(url)
            return parsed.netloc.lower()
        except:
            return url.lower()
    
    def add_to_whitelist(self, address: str, note: str = "") -> bool:
        """Add address to whitelist"""
        try:
            address = address.lower()
            self.whitelist.add(address)
            return True
        except:
            return False
    
    def add_to_blacklist(self, address: str, note: str = "") -> bool:
        """Add address to blacklist"""
        try:
            address = address.lower()
            self.blacklist.add(address)
            return True
        except:
            return False
    
    def remove_from_whitelist(self, address: str) -> bool:
        """Remove address from whitelist"""
        try:
            address = address.lower()
            self.whitelist.discard(address)
            return True
        except:
            return False
    
    def remove_from_blacklist(self, address: str) -> bool:
        """Remove address from blacklist"""
        try:
            address = address.lower()
            self.blacklist.discard(address)
            return True
        except:
            return False
    
    def is_whitelisted(self, address: str) -> bool:
        """Check if address is whitelisted"""
        return address.lower() in self.whitelist
    
    def is_blacklisted(self, address: str) -> bool:
        """Check if address is blacklisted"""
        return address.lower() in self.blacklist
    
    def get_security_score(self, address: str) -> Dict[str, Any]:
        """Get overall security score for address"""
        score = 50  # Base score
        factors = []
        
        if self.is_whitelisted(address):
            score += 30
            factors.append("Whitelisted address")
        
        if self.is_blacklisted(address):
            score -= 40
            factors.append("Blacklisted address")
        
        if address.lower() in self.scam_addresses:
            score -= 50
            factors.append("Known scam address")
        
        if address.lower() in self.known_contracts:
            contract_info = self.known_contracts[address.lower()]
            if contract_info.get('verified'):
                score += 15
                factors.append("Verified contract")
            if contract_info.get('risk_level') == 'low':
                score += 10
                factors.append("Low risk contract")
        
        # Clamp score between 0 and 100
        score = max(0, min(100, score))
        
        if score >= 80:
            risk_level = RiskLevel.LOW
        elif score >= 60:
            risk_level = RiskLevel.MEDIUM
        elif score >= 40:
            risk_level = RiskLevel.HIGH
        else:
            risk_level = RiskLevel.CRITICAL
        
        return {
            "score": score,
            "risk_level": risk_level.value,
            "factors": factors,
            "recommendation": self._get_recommendation_for_score(score)
        }
    
    def _get_recommendation_for_score(self, score: int) -> str:
        """Get recommendation based on security score"""
        if score >= 80:
            return "Safe to interact"
        elif score >= 60:
            return "Proceed with caution"
        elif score >= 40:
            return "High risk - verify carefully"
        else:
            return "Do not interact - very high risk"
    
    def generate_security_report(self, address: str) -> Dict[str, Any]:
        """Generate comprehensive security report"""
        report = {
            "address": address,
            "timestamp": time.time(),
            "security_score": self.get_security_score(address),
            "whitelist_status": self.is_whitelisted(address),
            "blacklist_status": self.is_blacklisted(address),
            "known_scam": address.lower() in self.scam_addresses,
            "contract_info": self.known_contracts.get(address.lower()),
            "recommendations": []
        }
        
        # Add recommendations based on findings
        if report["known_scam"]:
            report["recommendations"].append("CRITICAL: This is a known scam address")
        
        if report["blacklist_status"]:
            report["recommendations"].append("WARNING: Address is on your blacklist")
        
        if report["whitelist_status"]:
            report["recommendations"].append("SAFE: Address is on your whitelist")
        
        if not report["contract_info"] and address != "0x":
            report["recommendations"].append("INFO: Unknown contract - verify before interacting")
        
        return report
    
    def export_security_settings(self) -> str:
        """Export security settings to JSON"""
        settings = {
            "whitelist": list(self.whitelist),
            "blacklist": list(self.blacklist),
            "version": "1.0",
            "exported_at": time.time()
        }
        return json.dumps(settings, indent=2)
    
    def import_security_settings(self, json_data: str) -> bool:
        """Import security settings from JSON"""
        try:
            settings = json.loads(json_data)
            
            if "whitelist" in settings:
                self.whitelist.update(settings["whitelist"])
            
            if "blacklist" in settings:
                self.blacklist.update(settings["blacklist"])
            
            return True
            
        except Exception as e:
            print(f"Failed to import security settings: {e}")
            return False

# Example usage
if __name__ == "__main__":
    security_manager = SecurityManager()
    
    # Test transaction risk assessment
    test_transaction = {
        "from": "0x742d35Cc6639C0532fA5B3De4839C4C6cCE3c03D",
        "to": "0x000000000000000000000000000000000000dead",
        "value": "1000000000000000000",  # 1 ETH
        "gas": 21000,
        "gasPrice": "20000000000",  # 20 Gwei
        "data": "0x"
    }
    
    assessment = security_manager.assess_transaction_risk(test_transaction, None)
    print(f"Risk Level: {assessment.risk_level.value}")
    print(f"Risk Score: {assessment.score}")
    print(f"Alerts: {[alert.value for alert in assessment.alerts]}")
    print(f"Warnings: {assessment.warnings}")
    
    # Test website security check
    website_check = security_manager.check_website_security("unisvvap.org")
    print(f"\nWebsite Security: {website_check}")
    
    # Test security score
    security_score = security_manager.get_security_score("0x742d35Cc6639C0532fA5B3De4839C4C6cCE3c03D")
    print(f"\nSecurity Score: {security_score}")