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('###################')