import boto3
import argparse
import json

def audit_s3_public_access(
    region_name='us-east-1',
    tag_key='public',
    tag_value='true',
    apply_tag=False,
    report_file_path=None
):
    """
    Audits S3 buckets in a given region for public access and optionally applies a tag
    to publicly accessible buckets.

    Args:
        region_name (str): The AWS region to audit.
        tag_key (str): The tag key to apply to publicly accessible buckets.
        tag_value (str): The tag value to apply to publicly accessible buckets.
        apply_tag (bool): If True, applies the specified tag to publicly accessible buckets.
        report_file_path (str, optional): Path to a file to save the report. If None, prints to console.
    """
    s3_client = boto3.client('s3', region_name=region_name)

    public_buckets = []

    print(f"Starting S3 public access audit in region {region_name}...")

    try:
        # 1. List all buckets
        print("\n>>> Step 1: Listing all S3 buckets...")
        response = s3_client.list_buckets()
        buckets = response['Buckets']
        print(f"   Found {len(buckets)} S3 buckets.")

        # 2. Check each bucket's public access settings
        print("\n>>> Step 2: Checking public access settings for each bucket...")
        for bucket in buckets:
            bucket_name = bucket['Name']
            is_public = False
            public_reason = []

            # Check Block Public Access settings
            try:
                bpa_response = s3_client.get_public_access_block(Bucket=bucket_name)
                bpa_config = bpa_response['PublicAccessBlockConfiguration']
                if bpa_config['BlockPublicAcls'] or bpa_config['IgnorePublicAcls'] or \
                   bpa_config['BlockPublicPolicy'] or bpa_config['RestrictPublicBuckets']:
                    # If any BPA setting is true, it's generally not public via those means
                    pass
                else:
                    # If all BPA settings are false, it might be public
                    is_public = True
                    public_reason.append("Block Public Access settings are not fully enabled.")
            except s3_client.exceptions.ClientError as e:
                if e.response['Error']['Code'] == 'NoSuchPublicAccessBlockConfiguration':
                    is_public = True
                    public_reason.append("No Block Public Access configuration found.")
                else:
                    print(f"      Error getting BPA for {bucket_name}: {e}")

            # Check Bucket Policy
            try:
                policy_response = s3_client.get_bucket_policy(Bucket=bucket_name)
                policy_doc = json.loads(policy_response['Policy'])
                for statement in policy_doc.get('Statement', []):
                    if statement.get('Effect') == 'Allow' and \
                       statement.get('Principal') == '*' and \
                       's3:GetObject' in statement.get('Action', []):
                        is_public = True
                        public_reason.append("Bucket policy allows public read access.")
                        break
            except s3_client.exceptions.ClientError as e:
                if e.response['Error']['Code'] == 'NoSuchBucketPolicy':
                    pass # No policy, so not public via policy
                else:
                    print(f"      Error getting bucket policy for {bucket_name}: {e}")

            # Check Bucket ACLs (legacy, but still possible)
            try:
                acl_response = s3_client.get_bucket_acl(Bucket=bucket_name)
                for grant in acl_response['Grants']:
                    if grant.get('Grantee', {}).get('URI') == 'http://acs.amazonaws.com/groups/global/AllUsers' and \
                       grant.get('Permission') in ['READ', 'FULL_CONTROL']:
                        is_public = True
                        public_reason.append("Bucket ACL allows public read access.")
                        break
            except Exception as e:
                print(f"      Error getting bucket ACL for {bucket_name}: {e}")

            if is_public:
                public_buckets.append({
                    'BucketName': bucket_name,
                    'PublicReason': '; '.join(public_reason)
                })
                print(f"   Bucket '{bucket_name}' is PUBLIC. Reason: {'; '.join(public_reason)}")

                # 3. Optionally apply tag
                if apply_tag:
                    print(f"      Applying tag {tag_key}={tag_value} to bucket '{bucket_name}'...")
                    try:
                        s3_client.put_bucket_tagging(
                            Bucket=bucket_name,
                            Tagging={
                                'TagSet': [
                                    {'Key': tag_key, 'Value': tag_value}
                                ]
                            }
                        )
                        print(f"      Tag {tag_key}={tag_value} applied to '{bucket_name}'.")
                    except Exception as e:
                        print(f"      Error applying tag to {bucket_name}: {e}")

    except Exception as e:
        print(f"Error during S3 public access audit: {e}")
        return

    # 4. Generate Report
    print("\n>>> Step 4: Generating report...")
    report_output = []
    if public_buckets:
        report_output.append("--- S3 Buckets with Public Access ---")
        for bucket_info in public_buckets:
            report_output.append(f"  Bucket Name: {bucket_info['BucketName']}")
            report_output.append(f"  Reason: {bucket_info['PublicReason']}")
            report_output.append("----------------------------------------")
    else:
        report_output.append("No S3 buckets found with public access.")

    if report_file_path:
        with open(report_file_path, 'w') as f:
            for line in report_output:
                f.write(line + '\n')
        print(f"Report saved to '{report_file_path}'.")
    else:
        for line in report_output:
            print(line)

    print("\nS3 public access audit completed.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Audits S3 buckets for public access and optionally applies a tag.")
    parser.add_argument("--region", default="us-east-1", help="AWS region to audit (default: us-east-1).")
    parser.add_argument("--tag-key", default="public", help="Tag key to apply to publicly accessible buckets (default: 'public').")
    parser.add_argument("--tag-value", default="true", help="Tag value to apply to publicly accessible buckets (default: 'true').")
    parser.add_argument("--apply-tag", action="store_true", help="If set, applies the specified tag to publicly accessible buckets.")
    parser.add_argument("--report-file-path", help="Optional. Path to a file to save the report. If not provided, prints to console.")

    args = parser.parse_args()

    audit_s3_public_access(
        region_name=args.region,
        tag_key=args.tag_key,
        tag_value=args.tag_value,
        apply_tag=args.apply_tag,
        report_file_path=args.report_file_path
    )