import boto3
import os

"""
how to use this class

import aws_modules.get_all_sg_rules

sg_rule = aws_modules.get_all_sg_rules.sg(aws_account) # passing aws_account value to retrive all sg rules

sg_rule_result = sg_rule.getSgRules()

"""

class sg:
    def __init__(req, aws_account):
        req.aws_account = aws_account
    
    def getSgRules(req):
    
      try:

        os.environ["AWS_PROFILE"] = req.aws_account
        aws_client = boto3.client('ec2')
        paginator = aws_client.get_paginator('describe_security_groups')
        response_iterator = paginator.paginate(
            DryRun=False,
        )
        
        vpc_id = ""
        sg_name = ""
        sg_id = ""
        ip_permissions = ""
        from_port = ""
        to_port = ""
        protocol = ""
        ip_range = ""
        ip_ranges_description = ""
        user_id_sg = ""
        user_id_sg_description = ""
        row = ""
        result = []

    
        for page in response_iterator:
            for sg in page["SecurityGroups"]:
                sg_name = sg["GroupName"]
                vpc_id = sg["VpcId"]
                sg_id = sg["GroupId"]
                ip_permissions = sg["IpPermissions"]
                for perm in ip_permissions:
                    if "FromPort" in perm.keys():
                        from_port = perm["FromPort"]
                    else:
                        from_port = ""
                    if "ToPort" in perm.keys():
                        to_port = perm["ToPort"]
                    else:
                        to_port = ""
                    protocol = perm["IpProtocol"]
                    for ip_ranges in perm["IpRanges"]:
                        ip_range = ip_ranges["CidrIp"]
                        if "Description" in ip_ranges.keys():
                            ip_ranges_description = ip_ranges["Description"]
                        else:
                         ip_ranges_description = ""
                        if  ip_range:
                            row = [vpc_id, sg_name, sg_id, from_port, to_port, protocol, ip_range, ip_ranges_description]
                            result.append(row)
                    for userid in perm["UserIdGroupPairs"]:
                        user_id_sg = userid["GroupId"]
                        if "Description" in userid.keys():
                            user_id_sg_description = userid["Description"]
                        else:
                            user_id_sg_description = ""
                        if  user_id_sg:
                            row = [vpc_id, sg_name, sg_id, from_port, to_port, protocol, user_id_sg, ip_ranges_description]
                            result.append(row)
        
        return result 

      except Exception as e:
          print(e)       
                

# boto3 is taking your aws creds for authorization 

import csv
import aws_modules.get_all_sg_rules

def main():

    aws_accounts = [
        "my-aws-account-1",
        "my-aws-account-2",
        "my-aws-account-3"
    ]

    try:
        for aws_account in aws_accounts:
             f = open("{0}_sg_rules.csv".format(aws_account), "w")
             writer = csv.DictWriter(
                 f, fieldnames=["vpc_id", "sg_name", "sg_id", "from port", "to port", "portocl", "source", "description"])
             writer.writeheader()
             f.close()

             sg_rule = aws_modules.get_all_sg_rules.sg(aws_account)
             sg_rule_result = sg_rule.getSgRules()

             for rule in sg_rule_result:
                 print("writing:", rule)
                 with open("{0}_sg_rules.csv".format(aws_account),"a") as fd:
                    wr = csv.writer(fd, dialect='excel')
                    wr.writerow(rule)

    except Exception as e:
          print(e) 


if __name__ == '__main__':
    main()