security_group_analyzer.py

import asyncio
import boto3
import pandas as pd
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, asdict
from datetime import datetime
import json
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

@dataclass
class SecurityGroupRule:
    """Data class for security group rule information"""
    account_id: str
    region: str
    security_group_id: str
    security_group_name: str
    rule_type: str  # 'ingress' or 'egress'
    ip_protocol: str
    from_port: Optional[int]
    to_port: Optional[int]
    cidr_blocks: List[str]
    security_group_refs: List[str]
    description: str
    is_public: bool
    vpc_id: str

class SecurityGroupAnalyzer:
    """Main class for analyzing AWS Security Groups across multiple accounts"""
    
    def __init__(self, account_ids: List[str], regions: Optional[List[str]] = None):
        """
        Initialize the analyzer with account IDs and regions
        
        Args:
            account_ids: List of AWS account IDs to analyze
            regions: List of AWS regions to analyze (defaults to common regions)
        """
        self.account_ids = account_ids
        self.regions = regions or [
            'us-east-1', 'us-west-2', 'eu-west-1', 'eu-central-1', 
            'ap-southeast-1', 'ap-northeast-1'
        ]
        self.security_groups_data: List[SecurityGroupRule] = []
        
    def get_session_for_account(self, account_id: str) -> boto3.Session:
        """
        Get AWS session for a specific account
        This assumes cross-account roles are set up
        """
        try:
            # Assume role in target account
            sts_client = boto3.client('sts')
            role_arn = f"arn:aws:iam::{account_id}:role/SecurityAuditRole"
            
            assumed_role = sts_client.assume_role(
                RoleArn=role_arn,
                RoleSessionName=f"security-group-analysis-{account_id}"
            )
            
            credentials = assumed_role['Credentials']
            return boto3.Session(
                aws_access_key_id=credentials['AccessKeyId'],
                aws_secret_access_key=credentials['SecretAccessKey'],
                aws_session_token=credentials['SessionToken']
            )
        except Exception as e:
            logger.error(f"Failed to assume role for account {account_id}: {e}")
            # Fallback to default session (useful for single account or local testing)
            return boto3.Session()
    
    def analyze_security_group_rules(self, ec2_client, account_id: str, region: str, sg: Dict) -> List[SecurityGroupRule]:
        """Analyze individual security group and extract rule information"""
        rules = []
        
        for rule_type, rule_list in [('ingress', sg.get('IpPermissions', [])), ('egress', sg.get('IpPermissionsEgress', []))]:
            for rule in rule_list:
                # Extract CIDR blocks
                cidr_blocks = [ip_range['CidrIp'] for ip_range in rule.get('IpRanges', [])]
                
                # Extract security group references
                sg_refs = [ref['GroupId'] for ref in rule.get('UserIdGroupPairs', [])]
                
                # Check if rule allows public access (0.0.0.0/0)
                is_public = '0.0.0.0/0' in cidr_blocks
                
                # Create rule object
                sg_rule = SecurityGroupRule(
                    account_id=account_id,
                    region=region,
                    security_group_id=sg['GroupId'],
                    security_group_name=sg.get('GroupName', ''),
                    rule_type=rule_type,
                    ip_protocol=rule.get('IpProtocol', ''),
                    from_port=rule.get('FromPort'),
                    to_port=rule.get('ToPort'),
                    cidr_blocks=cidr_blocks,
                    security_group_refs=sg_refs,
                    description=rule.get('Description', ''),
                    is_public=is_public,
                    vpc_id=sg.get('VpcId', '')
                )
                rules.append(sg_rule)
        
        return rules
    
    def collect_security_groups_for_account_region(self, account_id: str, region: str) -> List[SecurityGroupRule]:
        """Collect security group data for a specific account and region"""
        logger.info(f"Analyzing account {account_id} in region {region}")
        
        try:
            session = self.get_session_for_account(account_id)
            ec2_client = session.client('ec2', region_name=region)
            
            response = ec2_client.describe_security_groups()
            security_groups = response['SecurityGroups']
            
            all_rules = []
            for sg in security_groups:
                rules = self.analyze_security_group_rules(ec2_client, account_id, region, sg)
                all_rules.extend(rules)
            
            logger.info(f"Found {len(all_rules)} rules in account {account_id}, region {region}")
            return all_rules
            
        except Exception as e:
            logger.error(f"Error analyzing account {account_id} in region {region}: {e}")
            return []
    
    def collect_all_security_groups(self) -> None:
        """Collect security group data from all accounts and regions using parallel processing"""
        logger.info("Starting security group data collection across all accounts and regions")
        
        all_rules = []
        
        # Use ThreadPoolExecutor for parallel processing
        with ThreadPoolExecutor(max_workers=10) as executor:
            futures = []
            
            for account_id in self.account_ids:
                for region in self.regions:
                    future = executor.submit(
                        self.collect_security_groups_for_account_region, 
                        account_id, 
                        region
                    )
                    futures.append(future)
            
            # Collect results as they complete
            for future in as_completed(futures):
                try:
                    rules = future.result()
                    all_rules.extend(rules)
                except Exception as e:
                    logger.error(f"Error in thread execution: {e}")
        
        self.security_groups_data = all_rules
        logger.info(f"Collected total of {len(all_rules)} security group rules")
    
    def get_public_rules(self) -> pd.DataFrame:
        """Get all rules that allow public access (0.0.0.0/0)"""
        public_rules = [rule for rule in self.security_groups_data if rule.is_public]
        return pd.DataFrame([asdict(rule) for rule in public_rules])
    
    def get_communication_map(self) -> pd.DataFrame:
        """Create a communication map between security groups"""
        communications = []
        
        for rule in self.security_groups_data:
            for sg_ref in rule.security_group_refs:
                communications.append({
                    'source_account': rule.account_id,
                    'source_region': rule.region,
                    'source_sg_id': rule.security_group_id,
                    'source_sg_name': rule.security_group_name,
                    'target_sg_id': sg_ref,
                    'rule_type': rule.rule_type,
                    'protocol': rule.ip_protocol,
                    'from_port': rule.from_port,
                    'to_port': rule.to_port,
                    'description': rule.description
                })
        
        return pd.DataFrame(communications)
    
    def get_port_analysis(self) -> pd.DataFrame:
        """Analyze port usage across all security groups"""
        port_data = []
        
        for rule in self.security_groups_data:
            if rule.from_port is not None:
                port_data.append({
                    'account_id': rule.account_id,
                    'region': rule.region,
                    'security_group_id': rule.security_group_id,
                    'rule_type': rule.rule_type,
                    'protocol': rule.ip_protocol,
                    'port': rule.from_port if rule.from_port == rule.to_port else f"{rule.from_port}-{rule.to_port}",
                    'is_public': rule.is_public,
                    'cidr_count': len(rule.cidr_blocks)
                })
        
        return pd.DataFrame(port_data)
    
    def export_to_csv(self, output_dir: str = "./reports") -> None:
        """Export analysis results to CSV files"""
        import os
        os.makedirs(output_dir, exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Export all security group rules
        all_rules_df = pd.DataFrame([asdict(rule) for rule in self.security_groups_data])
        all_rules_df.to_csv(f"{output_dir}/all_security_group_rules_{timestamp}.csv", index=False)
        
        # Export public rules
        public_rules_df = self.get_public_rules()
        public_rules_df.to_csv(f"{output_dir}/public_rules_{timestamp}.csv", index=False)
        
        # Export communication map
        comm_map_df = self.get_communication_map()
        comm_map_df.to_csv(f"{output_dir}/communication_map_{timestamp}.csv", index=False)
        
        # Export port analysis
        port_analysis_df = self.get_port_analysis()
        port_analysis_df.to_csv(f"{output_dir}/port_analysis_{timestamp}.csv", index=False)
        
        logger.info(f"Reports exported to {output_dir}")
    
    def get_summary_stats(self) -> Dict[str, Any]:
        """Get summary statistics of the security group analysis"""
        total_rules = len(self.security_groups_data)
        public_rules = len([rule for rule in self.security_groups_data if rule.is_public])
        
        accounts_analyzed = len(set(rule.account_id for rule in self.security_groups_data))
        regions_analyzed = len(set(rule.region for rule in self.security_groups_data))
        unique_sgs = len(set(rule.security_group_id for rule in self.security_groups_data))
        
        return {
            'total_rules': total_rules,
            'public_rules': public_rules,
            'public_percentage': (public_rules / total_rules * 100) if total_rules > 0 else 0,
            'accounts_analyzed': accounts_analyzed,
            'regions_analyzed': regions_analyzed,
            'unique_security_groups': unique_sgs,
            'analysis_timestamp': datetime.now().isoformat()
        }


def main():
    """Main function to run the security group analysis"""
    # Define your AWS account IDs
    account_ids = [
        '549166761550', '025381531841', '138351723750', '852584260383',
        '790679477109', '144196841275', '174436502675', '575045511530',
        '713037428386', '856231879353', '338589391164', '147367150109',
        '270134153674', '106363003370'
    ]
    
    # Initialize analyzer
    analyzer = SecurityGroupAnalyzer(account_ids)
    
    # Collect data
    analyzer.collect_all_security_groups()
    
    # Export results
    analyzer.export_to_csv()
    
    # Print summary
    summary = analyzer.get_summary_stats()
    print("\n=== Security Group Analysis Summary ===")
    print(f"Total Rules Analyzed: {summary['total_rules']}")
    print(f"Public Rules (0.0.0.0/0): {summary['public_rules']} ({summary['public_percentage']:.1f}%)")
    print(f"Accounts Analyzed: {summary['accounts_analyzed']}")
    print(f"Regions Analyzed: {summary['regions_analyzed']}")
    print(f"Unique Security Groups: {summary['unique_security_groups']}")
    print(f"Analysis completed at: {summary['analysis_timestamp']}")


if __name__ == "__main__":
    main()