#!/usr/bin/env python3
"""
GlassWorm VS Code Extension Scanner
Detects and removes malicious VS Code extensions infected by the GlassWorm worm
Based on Koi Security research (October 2025)
"""

import os
import sys
import json
import shutil
import platform
import subprocess
from pathlib import Path
from typing import List, Dict, Set, Tuple
import argparse


class GlassWormScanner:
    def __init__(self):
        # Known infected extensions from Koi Security research
        self.infected_extensions = {
            "codejoy.codejoy-vscode-extension": ["1.8.3", "1.8.4"],
            "l-igh-t.vscode-theme-seti-folder": ["1.2.3"],
            "kleinesfilmroellchen.serenity-dsl-syntaxhighlight": ["0.3.2"],
            "JScearcy.rust-doc-viewer": ["4.2.1"],
            "SIRILMP.dark-theme-sm": ["3.11.4"],
            "CodeInKlingon.git-worktree-menu": ["1.0.9", "1.0.91"],
            "ginfuru.better-nunjucks": ["0.3.2"],
            "ellacrity.recoil": ["0.7.4"],
            "grrrck.positron-plus-1-e": ["0.0.71"],
            "jeronimoekerdt.color-picker-universal": ["2.8.91"],
            "srcery-colors.srcery-colors": ["0.3.9"],
            "sissel.shopify-liquid": ["4.0.1"],
            "TretinV3.forts-api-extention": ["0.3.1"],
            "cline-ai-main.cline-ai-agent": ["3.1.3"]
        }
        
        # Known malicious indicators
        self.c2_servers = [
            "217.69.3.218",
            "199.247.10.166",
            "140.82.52.31"
        ]
        
        self.solana_wallet = "28PKnu7RzizxBzFPoLp69HLXp9bJL3JFtT2s5QzHsEA2"
        
        # Unicode variation selectors (invisible characters)
        self.variation_selectors = [
            "\uFE00", "\uFE01", "\uFE02", "\uFE03", "\uFE04", "\uFE05",
            "\uFE06", "\uFE07", "\uFE08", "\uFE09", "\uFE0A", "\uFE0B",
            "\uFE0C", "\uFE0D", "\uFE0E", "\uFE0F"
        ]
        
        self.findings: List[Dict] = []
        self.removed_extensions: List[str] = []

    def get_vscode_extensions_path(self) -> List[Path]:
        """Get VS Code extensions directory based on OS"""
        system = platform.system()
        paths = []
        
        if system == "Windows":
            user_profile = os.getenv('USERPROFILE')
            if user_profile:
                paths.append(Path(user_profile) / ".vscode" / "extensions")
                paths.append(Path(user_profile) / ".vscode-insiders" / "extensions")
        elif system == "Darwin":  # macOS
            home = Path.home()
            paths.append(home / ".vscode" / "extensions")
            paths.append(home / ".vscode-insiders" / "extensions")
        elif system == "Linux":
            home = Path.home()
            paths.append(home / ".vscode" / "extensions")
            paths.append(home / ".vscode-insiders" / "extensions")
            paths.append(home / ".vscode-oss" / "extensions")
        
        return [p for p in paths if p.exists()]

    def get_installed_extensions(self) -> List[Tuple[str, str, Path]]:
        """Get list of installed VS Code extensions"""
        extensions = []
        
        for ext_dir in self.get_vscode_extensions_path():
            if not ext_dir.exists():
                continue
            
            for item in ext_dir.iterdir():
                if not item.is_dir():
                    continue
                
                # Extension folder format: publisher.name-version
                dir_name = item.name
                
                # Parse extension ID and version
                if '-' in dir_name:
                    parts = dir_name.rsplit('-', 1)
                    ext_id = parts[0]
                    version = parts[1] if len(parts) > 1 else "unknown"
                else:
                    ext_id = dir_name
                    version = "unknown"
                
                # Try to get more accurate info from package.json
                package_json = item / "package.json"
                if package_json.exists():
                    try:
                        with open(package_json, 'r', encoding='utf-8') as f:
                            data = json.load(f)
                            publisher = data.get('publisher', '')
                            name = data.get('name', '')
                            version = data.get('version', version)
                            
                            if publisher and name:
                                ext_id = f"{publisher}.{name}"
                    except:
                        pass
                
                extensions.append((ext_id, version, item))
        
        return extensions

    def check_for_invisible_unicode(self, file_path: Path) -> List[str]:
        """Check file for invisible Unicode variation selectors"""
        found_selectors = []
        
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                
                for selector in self.variation_selectors:
                    if selector in content:
                        found_selectors.append(hex(ord(selector)))
        except:
            pass
        
        return found_selectors

    def check_for_malicious_patterns(self, ext_path: Path) -> Dict:
        """Check extension for malicious patterns"""
        indicators = {
            'c2_servers': [],
            'solana_references': False,
            'invisible_unicode': [],
            'suspicious_files': [],
            'google_calendar_c2': False,
            'credential_theft': False
        }
        
        # Search through JS/TS files
        for file_path in ext_path.rglob('*.js'):
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()
                    
                    # Check for C2 servers
                    for c2 in self.c2_servers:
                        if c2 in content:
                            indicators['c2_servers'].append(c2)
                    
                    # Check for Solana wallet
                    if self.solana_wallet in content:
                        indicators['solana_references'] = True
                    
                    # Check for invisible Unicode
                    unicode_found = self.check_for_invisible_unicode(file_path)
                    if unicode_found:
                        indicators['invisible_unicode'].extend(unicode_found)
                        indicators['suspicious_files'].append(str(file_path))
                    
                    # Check for Google Calendar C2
                    if 'calendar.google.com' in content or 'googleapis.com/calendar' in content:
                        indicators['google_calendar_c2'] = True
                    
                    # Check for credential theft indicators
                    credential_keywords = ['npm', 'token', 'github', 'git credential', '.npmrc', 'vsix']
                    if any(keyword in content.lower() for keyword in credential_keywords):
                        indicators['credential_theft'] = True
                    
            except:
                pass
        
        return indicators

    def scan_extensions(self, remove_infected: bool = False):
        """Scan all VS Code extensions for GlassWorm infection"""
        print("="*70)
        print("GlassWorm VS Code Extension Scanner")
        print("Detecting malicious extensions from October 2025 supply chain attack")
        print("="*70)
        print()
        
        extensions = self.get_installed_extensions()
        
        if not extensions:
            print("❌ No VS Code extensions found!")
            return
        
        print(f"📦 Found {len(extensions)} installed extensions")
        print()
        
        infected_count = 0
        suspicious_count = 0
        
        for ext_id, version, ext_path in extensions:
            is_known_infected = False
            
            # Check if it's a known infected extension
            if ext_id in self.infected_extensions:
                if version in self.infected_extensions[ext_id] or version == "unknown":
                    is_known_infected = True
                    infected_count += 1
                    
                    print(f"🚨 CRITICAL: Known infected extension found!")
                    print(f"   Extension: {ext_id}")
                    print(f"   Version: {version}")
                    print(f"   Location: {ext_path}")
                    print()
                    
                    self.findings.append({
                        'id': ext_id,
                        'version': version,
                        'path': str(ext_path),
                        'status': 'KNOWN_INFECTED',
                        'risk_score': 100
                    })
                    
                    if remove_infected:
                        self.remove_extension(ext_path, ext_id)
            
            # Deep scan for malicious patterns
            if not is_known_infected:
                indicators = self.check_for_malicious_patterns(ext_path)
                
                risk_score = 0
                reasons = []
                
                if indicators['c2_servers']:
                    risk_score += 50
                    reasons.append(f"C2 servers found: {', '.join(indicators['c2_servers'])}")
                
                if indicators['solana_references']:
                    risk_score += 30
                    reasons.append("Solana blockchain references")
                
                if indicators['invisible_unicode']:
                    risk_score += 40
                    reasons.append(f"Invisible Unicode characters: {', '.join(set(indicators['invisible_unicode']))}")
                
                if indicators['google_calendar_c2']:
                    risk_score += 20
                    reasons.append("Google Calendar API usage (potential C2)")
                
                if indicators['credential_theft']:
                    risk_score += 15
                    reasons.append("Credential access patterns")
                
                if risk_score >= 50:
                    suspicious_count += 1
                    status = "🚨 HIGH RISK" if risk_score >= 80 else "⚠️  SUSPICIOUS"
                    
                    print(f"{status}")
                    print(f"   Extension: {ext_id}")
                    print(f"   Version: {version}")
                    print(f"   Risk Score: {risk_score}")
                    print(f"   Location: {ext_path}")
                    print(f"   Indicators:")
                    for reason in reasons:
                        print(f"      • {reason}")
                    print()
                    
                    self.findings.append({
                        'id': ext_id,
                        'version': version,
                        'path': str(ext_path),
                        'status': 'SUSPICIOUS',
                        'risk_score': risk_score,
                        'indicators': reasons
                    })
                    
                    if remove_infected and risk_score >= 80:
                        self.remove_extension(ext_path, ext_id)
        
        print("="*70)
        print("📊 SCAN SUMMARY")
        print("="*70)
        print(f"Total Extensions Scanned: {len(extensions)}")
        print(f"🚨 Known Infected: {infected_count}")
        print(f"⚠️  Suspicious: {suspicious_count}")
        print(f"✅ Clean: {len(extensions) - infected_count - suspicious_count}")
        
        if self.removed_extensions:
            print(f"\n🗑️  Removed Extensions: {len(self.removed_extensions)}")
            for ext in self.removed_extensions:
                print(f"   • {ext}")

    def remove_extension(self, ext_path: Path, ext_id: str):
        """Remove an infected extension"""
        try:
            print(f"   🗑️  Removing extension: {ext_id}")
            shutil.rmtree(ext_path)
            self.removed_extensions.append(ext_id)
            print(f"   ✅ Successfully removed!")
        except Exception as e:
            print(f"   ❌ Failed to remove: {e}")
        print()

    def check_system_compromise(self):
        """Check for system-level compromise indicators"""
        print("\n" + "="*70)
        print("🔍 CHECKING FOR SYSTEM COMPROMISE")
        print("="*70)
        
        # Check for SOCKS proxy
        print("\n[1] Checking for unauthorized SOCKS proxy...")
        # This would require more system-specific checks
        
        # Check for suspicious network connections
        print("[2] Checking for suspicious network connections...")
        suspicious_found = False
        
        try:
            if platform.system() == "Windows":
                result = subprocess.run(['netstat', '-ano'], capture_output=True, text=True)
            else:
                result = subprocess.run(['netstat', '-tuln'], capture_output=True, text=True)
            
            output = result.stdout
            
            for c2 in self.c2_servers:
                if c2 in output:
                    print(f"   🚨 SUSPICIOUS: Connection to known C2 server: {c2}")
                    suspicious_found = True
        except:
            print("   ℹ️  Unable to check network connections")
        
        if not suspicious_found:
            print("   ✅ No suspicious network connections detected")
        
        # Check for credential files
        print("\n[3] Checking for potential credential theft...")
        home = Path.home()
        
        sensitive_files = [
            home / ".npmrc",
            home / ".gitconfig",
            home / ".git-credentials"
        ]
        
        for file_path in sensitive_files:
            if file_path.exists():
                print(f"   ⚠️  Found: {file_path}")
                print(f"      Recommendation: Review and rotate credentials")

    def generate_report(self, output_file: str = "glassworm_scan_report.json"):
        """Generate detailed JSON report"""
        if self.findings:
            with open(output_file, 'w') as f:
                json.dump(self.findings, f, indent=2)
            print(f"\n💾 Detailed report saved to: {output_file}")

    def print_recommendations(self):
        """Print security recommendations"""
        print("\n" + "="*70)
        print("🛡️  SECURITY RECOMMENDATIONS")
        print("="*70)
        print("""
1. IMMEDIATE ACTIONS:
   • Remove all infected extensions immediately
   • Restart VS Code completely
   • Check VS Code extension auto-update settings
   
2. CREDENTIAL SECURITY:
   • Rotate ALL credentials (GitHub, npm, Git, API tokens)
   • Review GitHub OAuth apps and revoke suspicious ones
   • Check npm access tokens: npm token list
   • Review .npmrc and .git-credentials files
   
3. CRYPTOCURRENCY WALLETS:
   • Check all cryptocurrency wallet extensions
   • Review transaction history for unauthorized transfers
   • Move funds to new wallets with new keys
   
4. SYSTEM SECURITY:
   • Scan for SOCKS proxy or VNC servers
   • Monitor network traffic for suspicious connections
   • Check running processes for unusual activity
   • Run full antivirus/antimalware scan
   
5. PREVENTION:
   • Only install extensions from trusted publishers
   • Disable automatic extension updates
   • Regularly audit installed extensions
   • Use extension allow-lists in enterprise environments
   
6. MONITORING:
   • Watch for unauthorized code commits to repositories
   • Monitor npm package publishing activity
   • Check for unexpected VS Code extension updates
   
For more information:
https://www.koi.ai/blog/glassworm-first-self-propagating-worm
        """)


def main():
    parser = argparse.ArgumentParser(
        description='Scan VS Code extensions for GlassWorm malware'
    )
    parser.add_argument(
        '--remove',
        action='store_true',
        help='Automatically remove infected extensions'
    )
    parser.add_argument(
        '--system-check',
        action='store_true',
        help='Check for system-level compromise'
    )
    parser.add_argument(
        '--report',
        help='Output report file path',
        default='glassworm_scan_report.json'
    )
    
    args = parser.parse_args()
    
    scanner = GlassWormScanner()
    scanner.scan_extensions(remove_infected=args.remove)
    
    if args.system_check:
        scanner.check_system_compromise()
    
    if scanner.findings:
        scanner.generate_report(args.report)
    
    scanner.print_recommendations()


if __name__ == "__main__":
    main()