from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class WellArchitectedAgent(BaseAgent):
    """
    An agent specialized in providing recommendations based on the AWS Well-Architected Framework.
    Starting with security posture analysis.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Well-Architected analysis.
        """
        if command == 'scan_security_posture':
            return self._scan_security_posture()
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by WellArchitectedAgent.")

    def _scan_security_posture(self):
        """
        Scans the AWS account for common security best practice violations.
        Checks for:
        - Publicly accessible S3 buckets.
        - EC2 Security Groups with unrestricted (0.0.0.0/0) SSH/RDP access.
        - IAM users with AdministratorAccess policy directly attached.
        """
        print("WellArchitectedAgent: Scanning account for security posture violations...")
        findings = []

        try:
            # Get all available regions for EC2/S3 checks
            ec2_client_global = AWSConnector.get_client('ec2', region_name='us-east-1') # Use a global client for region list
            regions = [region['RegionName'] for region in ec2_client_global.describe_regions()['Regions']]

            # --- Check S3 Buckets for Public Access ---
            s3_client = AWSConnector.get_client('s3')
            response = s3_client.list_buckets()
            for bucket in response.get('Buckets', []):
                bucket_name = bucket['Name']
                try:
                    # Check Block Public Access settings
                    bpa_status = s3_client.get_public_access_block(Bucket=bucket_name)
                    if not bpa_status['PublicAccessBlockConfiguration']['BlockPublicAcls'] or \
                       not bpa_status['PublicAccessBlockConfiguration']['IgnorePublicAcls'] or \
                       not bpa_status['PublicAccessBlockConfiguration']['BlockPublicPolicy'] or \
                       not bpa_status['PublicAccessBlockConfiguration']['RestrictPublicBuckets']:
                        findings.append(f"S3 Security: Bucket '{bucket_name}' does NOT have all Block Public Access settings enabled. Review for public access.")
                except s3_client.exceptions.NoSuchPublicAccessBlockConfiguration:
                    findings.append(f"S3 Security: Bucket '{bucket_name}' does NOT have Block Public Access configured. It might be publicly accessible.")
                except Exception as e:
                    # Handle cases where we might not have permission to check BPA
                    print(f"Warning: Could not check BPA for bucket {bucket_name}: {e}")

            # --- Check EC2 Security Groups for Unrestricted SSH/RDP ---
            for region in regions:
                region_ec2_client = AWSConnector.get_client('ec2', region_name=region)
                sgs = region_ec2_client.describe_security_groups().get('SecurityGroups', [])
                for sg in sgs:
                    for ip_permission in sg.get('IpPermissions', []):
                        for ip_range in ip_permission.get('IpRanges', []):
                            if ip_range['CidrIp'] == '0.0.0.0/0':
                                from_port = ip_permission.get('FromPort')
                                to_port = ip_permission.get('ToPort')
                                if from_port is not None and to_port is not None:
                                    if (from_port <= 22 <= to_port) or (from_port <= 3389 <= to_port):
                                        findings.append(f"EC2 Security: Security Group '{sg['GroupName']}' ({sg['GroupId']}) in region {region} allows unrestricted access (0.0.0.0/0) to port(s) {from_port}-{to_port}. This includes sensitive ports like SSH/RDP.")

            # --- Check IAM Users for AdministratorAccess ---
            iam_client = AWSConnector.get_client('iam')
            users = iam_client.list_users().get('Users', [])
            for user in users:
                user_name = user['UserName']
                attached_policies = iam_client.list_attached_user_policies(UserName=user_name).get('AttachedPolicies', [])
                for policy in attached_policies:
                    if policy['PolicyName'] == 'AdministratorAccess':
                        findings.append(f"IAM Security: User '{user_name}' has the 'AdministratorAccess' policy directly attached. Consider using more granular permissions.")

            if not findings:
                return {"status": "success", "message": "No immediate security posture findings for public S3 buckets, unrestricted SG access, or AdministratorAccess IAM users."}

            return {"status": "success", "findings": findings}
        except Exception as e:
            print(f"Error during security posture scan: {e}")
            return {"status": "error", "message": str(e)}