refactor code

Author: varendra007

.gitignore

@@ -11,4 +11,4 @@ crudActions/
 finalUserPolicies/
 userPolicies*/
 presentPolicies*/
-logs*/
+logs*/

CONTRIBUTING.md

@@ -1,4 +1,5 @@
 # Contributing to AWSZeroTrustPolicy
+
 We love your input! We want to make contributing to this project as easy and transparent as possible, whether it's:
 
 - Reporting a bug
@@ -8,9 +9,11 @@ We love your input! We want to make contributing to this project as easy and tra
 - Becoming a maintainer
 
 ## We Develop with Github
+
 We use github to host code, to track issues and feature requests, as well as accept pull requests.
 
 ## We Use [Github Flow](https://docs.github.com/en/get-started/quickstart/github-flow), So All Code Changes Happen Through Pull Requests
+
 Pull requests are the best way to propose changes to the codebase (we use [Github Flow](https://docs.github.com/en/get-started/quickstart/github-flow)). We actively welcome your pull requests:
 
 1. Fork the repo and create your branch from `master`.
@@ -21,14 +24,15 @@ Pull requests are the best way to propose changes to the codebase (we use [Githu
 6. Issue that pull request!
 
 ## Any contributions you make will be under the MIT Software License
+
 In short, when you submit code changes, your submissions are understood to be under the same [MIT License](http://choosealicense.com/licenses/mit/) that covers the project. Feel free to contact the maintainers if that's a concern.
 
 ## Report bugs using Github's [issues](https://github.com/CloudDefenseAI/AWSZeroTrustPolicy/issues)
+
 We use GitHub issues to track public bugs. Report a bug by [opening a new issue](https://github.com/CloudDefenseAI/AWSZeroTrustPolicy/issues/new); it's that easy!
 
 ## Write bug reports with detail, background, and sample code
 
-
 **Great Bug Reports** tend to have:
 
 - A quick summary and/or background
@@ -39,14 +43,15 @@ We use GitHub issues to track public bugs. Report a bug by [opening a new issue]
 - What actually happens
 - Notes (possibly including why you think this might be happening, or stuff you tried that didn't work)
 
-People *love* thorough bug reports. I'm not even kidding.
-
+People _love_ thorough bug reports. I'm not even kidding.
 
-* 2 spaces for indentation rather than tabs
-* You can try running `npm run lint` for style unification
+- 2 spaces for indentation rather than tabs
+- You can try running `npm run lint` for style unification
 
 ## License
+
 By contributing, you agree that your contributions will be licensed under its MIT License.
 
 ## References
+
 This document was adapted from the open-source contribution guidelines for [Facebook's Draft](https://github.com/facebook/draft-js/blob/a9316a723f9e918afde44dea68b5f9f39b7d9b00/CONTRIBUTING.md)

Dockerfile

@@ -6,6 +6,7 @@ RUN apk add --no-cache python3-dev gcc musl-dev libffi-dev libpq-dev openssl-dev
 WORKDIR /app
 COPY requirements.txt ./
 RUN pip install --no-cache-dir -r requirements.txt
+
 COPY . .
 EXPOSE 8000
 CMD ["sh", "-c", "redis-server & uvicorn app:app --reload --host 0.0.0.0"]

LICENSE

@@ -69,7 +69,7 @@ If you want to help and wish to contribute, please review our contribution guide
 
 ## License
 
-This project is released under the [Apache-2.0 License]([url](https://github.com/CloudDefenseAI/AWSZeroTrustPolicy/blob/master/LICENSE)).
+This project is released under the [Apache-2.0 License](<[url](https://github.com/CloudDefenseAI/AWSZeroTrustPolicy/blob/master/LICENSE)>).
 
 ## Disclaimer:
 

README.md

@@ -7,17 +7,30 @@
 
 app = FastAPI()
 
+
 @app.post("/run")
 def run_script(script: Script):
     print(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
     print("Payload:")
     print(json.dumps(script.dict(), indent=4))
 
     try:
-        resp = runner(script.accountType, script.accessKey, script.secretKey, script.accountId, script.days
-                      , script.bucketData, script.roleArn, script.externalId)
+        resp = runner(
+            script.accountType,
+            script.accessKey,
+            script.secretKey,
+            script.accountId,
+            script.days,
+            script.bucketData,
+            script.roleArn,
+            script.externalId,
+        )
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-    return {"accountId": resp['accountId'], "generatedPolicies": resp['generatedPolicies'], "consolidatedPolicies": resp['consolidatedPolicies'], "excessivePolicies": resp['excessivePolicies']}
-
+    return {
+        "accountId": resp["accountId"],
+        "generatedPolicies": resp["generatedPolicies"],
+        "consolidatedPolicies": resp["consolidatedPolicies"],
+        "excessivePolicies": resp["excessivePolicies"],
+    }

app.py

@@ -4,100 +4,110 @@
 from botocore.exceptions import ClientError, NoCredentialsError, BotoCoreError
 from fastapi import HTTPException
 
+
 class AWSOperations:
     def connect_to_iam_with_assumed_role(self, aws_credentials):
         # Create a new session with the temporary credentials
         session = boto3.Session(
             aws_access_key_id=aws_credentials.access_key,
             aws_secret_access_key=aws_credentials.secret_key,
-            aws_session_token=aws_credentials.token
+            aws_session_token=aws_credentials.token,
         )
         # Use the new session to connect to IAM
-        iam_client = session.client('iam')
+        iam_client = session.client("iam")
         return iam_client
 
     def get_iam_connection(self):
         try:
             with open("config.json", "r") as f:
                 data = json.loads(f.read())
-                if data['accountType'] == "CloudFormation":
+                if data["accountType"] == "CloudFormation":
                     aws_credentials = self.get_assume_role_credentials(data)
                     iam_client = self.connect_to_iam_with_assumed_role(aws_credentials)
-                elif data['accountType'] == "Credential":
-                    iam_client = self.connect("iam", data['aws_access_key_id'], data['aws_secret_access_key'])
+                elif data["accountType"] == "Credential":
+                    iam_client = self.connect(
+                        "iam", data["aws_access_key_id"], data["aws_secret_access_key"]
+                    )
                 return iam_client
         except (FileNotFoundError, json.JSONDecodeError) as e:
-            raise HTTPException(status_code=500, detail=f"Error reading or parsing config.json: {e}")
+            raise HTTPException(
+                status_code=500, detail=f"Error reading or parsing config.json: {e}"
+            )
 
         except (ClientError, NoCredentialsError, BotoCoreError) as e:
             raise HTTPException(status_code=500, detail=f"Error connecting to IAM: {e}")
 
-
-    def connect(self,serviceName,aws_access_key_id,aws_secret_access_key):
-        s3Client = boto3.client(serviceName,aws_access_key_id=aws_access_key_id,aws_secret_access_key=aws_secret_access_key)
+    def connect(self, serviceName, aws_access_key_id, aws_secret_access_key):
+        s3Client = boto3.client(
+            serviceName,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
         return s3Client
 
-    def connect_to_s3_with_assumed_role(self,aws_credentials):
+    def connect_to_s3_with_assumed_role(self, aws_credentials):
         # Create a new session with the temporary credentials
         session = boto3.Session(
             aws_access_key_id=aws_credentials.access_key,
             aws_secret_access_key=aws_credentials.secret_key,
-            aws_session_token=aws_credentials.token
+            aws_session_token=aws_credentials.token,
         )
         # Use the new session to connect to S3
-        s3_client = session.client('s3')
+        s3_client = session.client("s3")
         return s3_client
 
     def getConnection(self):
         try:
             with open("config.json", "r") as f:
                 data = json.loads(f.read())
-                if data['accountType'] == "CloudFormation":
+                if data["accountType"] == "CloudFormation":
                     aws_credentials = self.get_assume_role_credentials(data)
                     s3_client = self.connect_to_s3_with_assumed_role(aws_credentials)
-                elif data['accountType'] == "Credential":
-                    s3_client = self.connect("s3", data['aws_access_key_id'], data['aws_secret_access_key'])
+                elif data["accountType"] == "Credential":
+                    s3_client = self.connect(
+                        "s3", data["aws_access_key_id"], data["aws_secret_access_key"]
+                    )
                 return s3_client
         except (FileNotFoundError, json.JSONDecodeError) as e:
-            raise HTTPException(status_code=500, detail=f"Error reading or parsing config.json: {e}")
+            raise HTTPException(
+                status_code=500, detail=f"Error reading or parsing config.json: {e}"
+            )
         except (ClientError, NoCredentialsError, BotoCoreError) as e:
             raise HTTPException(status_code=500, detail=f"Error connecting to S3: {e}")
 
     def get_assume_role_credentials(self, account):
         try:
-        # Create an STS client with the IAM user's access and secret keys
+            # Create an STS client with the IAM user's access and secret keys
             sts_client = boto3.client(
-                'sts',
-                aws_access_key_id=account['aws_access_key_id'],
-                aws_secret_access_key=account['aws_secret_access_key']
+                "sts",
+                aws_access_key_id=account["aws_access_key_id"],
+                aws_secret_access_key=account["aws_secret_access_key"],
             )
 
             # Assume the IAM role
             response = sts_client.assume_role(
-                RoleArn=account['role_arn'],
-                RoleSessionName='Assume_Role_Session',
+                RoleArn=account["role_arn"],
+                RoleSessionName="Assume_Role_Session",
                 DurationSeconds=43200,
-                ExternalId=account['externalid']
+                ExternalId=account["externalid"],
             )
 
             # Extract the temporary credentials
-            creds = response['Credentials']
+            creds = response["Credentials"]
             session_credentials = boto3.Session(
-                aws_access_key_id=creds['AccessKeyId'],
-                aws_secret_access_key=creds['SecretAccessKey'],
-                aws_session_token=creds['SessionToken']
+                aws_access_key_id=creds["AccessKeyId"],
+                aws_secret_access_key=creds["SecretAccessKey"],
+                aws_session_token=creds["SessionToken"],
             ).get_credentials()
 
             # Create an AwsCredentials object with the temporary credentials
             aws_credentials = Credentials(
                 access_key=session_credentials.access_key,
                 secret_key=session_credentials.secret_key,
-                token=session_credentials.token
+                token=session_credentials.token,
             )
 
             return aws_credentials
 
         except (ClientError, NoCredentialsError, BotoCoreError) as e:
             raise HTTPException(status_code=500, detail=f"Error assuming role: {e}")
-
-

aws/awsOps.py

@@ -14,13 +14,15 @@
 # hdServices = ['a4b','account','amplify','apprunner','appsync','aps','billing','codebuild','codecommit','connect','databrew','eks','emrcontainers','forecast','frauddetector','fsx','gamelift','greengrassv2','health','iot','iotanalytics','iotevents','iotfleethub','iotthingsgraph','kafka','kendra','kinesisvideo','lakeformation','licensemanager','lookoutvision','macie','managedblockchain','marketplacecatalog','mediaconnect','mediaconvert','medialive','mediapackage','mediapackage-vod','mediastore','mediastore-data','mediatailor','meteringmarketplace','migrationhub-config','mobile','mq','neptune','networkmanager','outposts','personalize','pinpoint','pinpoint-email','pinpoint-sms-voice','polly','pricing','qldb','quicksight','ram','rds-data','robomaker','route53resolver','sagemaker','sagemaker-a2i-runtime','sagemaker-edge','sagemaker-featurestore-runtime','sagemaker-runtime','savingsplans','schemas','secretsmanager','securityhub','serverlessrepo','servicecatalog','servicecatalog-appregistry','servicequotas','sesv2','shield','signer','sms','snowball','snowball-edge','sso','sso-oidc','ssm','stepfunctions','storagegateway','synthetics','textract','transcribe','transfer','translate','waf-regional','wafv2','worklink','workmail','workmailmessageflow','workspaces','xray','autoscaling','iam','ec2','s3','rds','elasticache','elasticbeanstalk','elasticloadbalancing','elasticmapreduce','cloudfront','cloudtrail','cloudwatch','cloudwatchevents','cloudwatchlogs','config','datapipeline','directconnect','dynamodb','ecr','ecs','elasticfilesystem','elastictranscoder','glacier','kinesis','kms','lambda','opsworks','redshift','route53','route53domains','sdb','ses','sns','sqs','storagegateway','sts','support','swf','waf','workspaces','xray']
 # hdServices.extend(['acm','acm-pca','alexaforbusiness','amplifybackend','appconfig','appflow','appintegrations','appmesh','appstream','appsync','athena','auditmanager','autoscaling-plans','backup','batch','braket','budgets','ce','chime','cloud9','clouddirectory','cloudformation','cloudhsm','cloudhsmv2','cloudsearch','cloudsearchdomain','cloudtrail','cloudwatch','cloudwatchevents','cloudwatchlogs','codeartifact','codebuild','codecommit','codedeploy','codeguru-reviewer','codeguru-reviewer-runtime','codeguru-profiler','codeguru-profiler-runtime','codepipeline','codestar','codestar-connections','codestar-notifications','cognito-identity','cognito-idp','cognito-sync','comprehend','comprehendmedical','compute-optimizer','connect','connect-contact-lens','connectparticipant','cur','customer-profiles','dataexchange','datapipeline','datasync','dax','detective','devicefarm','devops-guru','directconnect','discovery','dlm','dms','docdb','ds','dynamodb','dynamodbstreams','ec2','ec2-instance-connect','ecr','ecr-public','ecs','eks','elastic-inference','elasticache','elasticbeanstalk','elasticfilesystem','elasticloadbalancing','elasticloadbalancingv2','elasticmapreduce','elastictranscoder','email','es','events','firehose','fms','forecast','forecastquery','frauddetector','fsx','gamelift','glacier','globalaccelerator','glue','greengrass','greengrassv2','groundstation','guardduty','health','healthlake','honeycode','iam','identitystore','imagebuilder','importexport','inspector','iot','iot-data','iot-jobs-data','iot1click-devices','iot1click-projects','iotanalytics','iotdeviceadvisor','iotevents','iotevents-data','iotfleethub','iotsecuretunneling','iotthingsgraph','iotwireless','ivs','kafka','kendra','kinesis','kinesis-video-archived-media','kinesis-video-media','kinesis-video-signaling','kinesisvideo','kinesisanalytics','kinesisanalyticsv2','kinesisvideoarchivedmedia','kinesis'])
 
+
 def create_services_list(actions_data):
     for action in actions_data:
         action_data = json.loads(action)
         event_source = action_data["eventSource"]
         service = event_source.split(".")[0]
         services[service].add(service)
 
+
 def create_service_actions_cache(services):
     service_actions_cache = {}
 
@@ -37,19 +39,25 @@ def create_service_actions_cache(services):
 
     return service_actions_cache
 
+
 def write_service_actions_cache_to_file(service_actions_cache, file_path):
-    with open(file_path, 'w') as f:
+    with open(file_path, "w") as f:
         json.dump(service_actions_cache, f, indent=2)
 
+
 def load_policy(filepath):
     with open(filepath, "r") as f:
         policy = json.load(f)
     return policy
 
+
 def is_valid_action(action):
-    return re.match(r'^[a-zA-Z0-9_]+:(\*|[a-zA-Z0-9_\*]+)$', action)
+    return re.match(r"^[a-zA-Z0-9_]+:(\*|[a-zA-Z0-9_\*]+)$", action)
 
-def compare_policy_worker(present_policy_filepath, user_policy_filepath, output_filepath):
+
+def compare_policy_worker(
+    present_policy_filepath, user_policy_filepath, output_filepath
+):
     print(f"Started thread for {user_policy_filepath}")
 
     current_policy = load_policy(present_policy_filepath)
@@ -59,7 +67,10 @@ def compare_policy_worker(present_policy_filepath, user_policy_filepath, output_
 
     with open(output_filepath, "w") as f_write:
         f_write.write(json.dumps(excessive_permissions, indent=2))
-        print(f"Generated excessive policy for {os.path.basename(user_policy_filepath)}")
+        print(
+            f"Generated excessive policy for {os.path.basename(user_policy_filepath)}"
+        )
+
 
 # def expand_wildcard_actions(actions_list, service_actions_cache=None):
 #     if service_actions_cache is None:
@@ -84,6 +95,7 @@ def compare_policy_worker(present_policy_filepath, user_policy_filepath, output_
 
 #     return expanded_actions
 
+
 def expand_wildcard_actions(actions_list, service_actions_cache=None):
     if service_actions_cache is None:
         with open("service_actions_cache.json", "r") as f:
@@ -98,35 +110,43 @@ def expand_wildcard_actions(actions_list, service_actions_cache=None):
         if is_valid_action(action):
             service, action_name = action.split(":")
             if "*" in action_name:
-                expanded_actions.extend([f"{a}" for a in service_actions_cache.get(service, []) if action_name.replace("*", "") in a])
+                expanded_actions.extend(
+                    [
+                        f"{a}"
+                        for a in service_actions_cache.get(service, [])
+                        if action_name.replace("*", "") in a
+                    ]
+                )
             else:
                 expanded_actions.append(action)
 
-        elif action == '*':
+        elif action == "*":
             for service in service_actions_cache:
-                expanded_actions.extend([f"{a}" for a in service_actions_cache[service]])
+                expanded_actions.extend(
+                    [f"{a}" for a in service_actions_cache[service]]
+                )
 
     return expanded_actions
 
+
 def compare_policy(current_policy, generated_policy):
-    excessive_permissions = {
-        "Version": "2012-10-17",
-        "Statement": []
-    }
+    excessive_permissions = {"Version": "2012-10-17", "Statement": []}
 
     for current_statement in current_policy["Statement"]:
         excessive_statement = {
             "Effect": current_statement["Effect"],
             "Action": [],
-            "Resource": current_statement["Resource"]
+            "Resource": current_statement["Resource"],
         }
 
         current_actions_expanded = expand_wildcard_actions(current_statement["Action"])
 
         for action in current_actions_expanded:
             action_in_generated = False
             for generated_statement in generated_policy["Statement"]:
-                generated_actions_expanded = expand_wildcard_actions(generated_statement["Action"])
+                generated_actions_expanded = expand_wildcard_actions(
+                    generated_statement["Action"]
+                )
 
                 # Check if the action and resource match in both policies
                 if action in generated_actions_expanded:
@@ -145,14 +165,17 @@ def compare_policy(current_policy, generated_policy):
 
     return excessive_permissions
 
+
 def compare_policies():
     crudKeys = crudConnection.get_all_keys()
     for user_arn in crudKeys:
         actions_data = crudConnection.get_list_items(user_arn)
         create_services_list(actions_data)
 
     service_actions_cache = create_service_actions_cache(services)
-    write_service_actions_cache_to_file(service_actions_cache, 'service_actions_cache.json')
+    write_service_actions_cache_to_file(
+        service_actions_cache, "service_actions_cache.json"
+    )
     print("Service actions cache created successfully.")
 
     # present_policies_dir = "presentPolicies"