ustayready · January 16, 2023 23:49 · Jan 16, 2023
diff --git a/gpt.py b/gpt.py
@@ -0,0 +1,73 @@
+import openai
+import boto3
+import json
+import time
+from typing import Dict, List
+
+openai.api_key = '### SET YOUR OPENAPI API KEY HERE ###'
+session = boto3.session.Session()
+client = session.client('iam')
+
+def get_role_names() -> List[str]:
+    """ Retrieve a list of role names by paginating over list_roles() calls """
+    roles = []
+    role_paginator = client.get_paginator('list_roles')
+    for response in role_paginator.paginate():
+        response_role_names = [r.get('RoleName') for r in response['Roles']]
+        roles.extend(response_role_names)
+    return roles
+
+def get_policies_for_roles(role_names: List[str]) -> Dict[str, List[Dict[str, str]]]:
+    """ Create a mapping of role names and any policies they have attached to them by
+        paginating over list_attached_role_policies() calls for each role name.
+        Attached policies will include policy name and ARN.
+    """
+    policy_map = {}
+    policy_paginator = client.get_paginator('list_attached_role_policies')
+    for name in role_names:
+        role_policies = []
+        for response in policy_paginator.paginate(RoleName=name):
+            role_policies.extend(response.get('AttachedPolicies'))
+        policy_map.update({name: role_policies})
+    return policy_map
+
+def check_policy(policy):
+    prompt = f'Does this AWS policy have any security vulnerabilities: \n{policy}'
+    response = openai.Completion.create(
+        model="text-davinci-003",
+        prompt=prompt,
+        temperature=0.5,
+        max_tokens=500,
+        top_p=1,
+        frequency_penalty=0.0,
+        presence_penalty=0.0,
+        stream=False,
+    )
+    answer = response.choices[0]['text']
+    print(answer)
+
+def retrieve_policy(arn):
+    policy = client.get_policy(
+        PolicyArn = arn
+    )
+    policy_version = client.get_policy_version(
+        PolicyArn = arn,
+        VersionId = policy['Policy']['DefaultVersionId']
+    )
+    return (policy, policy_version)
+
+role_names = get_role_names()
+attached_role_policies = get_policies_for_roles(role_names)
+
+for k, v in attached_role_policies.items():
+    for x in v:
+        name = k
+        arn = x['PolicyArn']
+        version, policy = retrieve_policy(arn)
+
+        print('###################')
+        print(f'{name} -> {arn}\n{policy}')
+        answer = check_policy(policy)
+
+        print(f'{answer}')
+        print('###################')
diff --git a/requirements.txt b/requirements.txt
@@ -0,0 +1,2 @@
+boto3
+openai