#!/usr/bin/python3
import sys
import boto3
import argparse

parser = argparse.ArgumentParser(description="Get SSH keys from IAM users")
parser.add_argument(
    "--user", dest="user", action="store", required=True, help="requested user"
)
parser.add_argument(
    "--group", action="store", required=True, help="IAM group to search"
)
parser.add_argument(
    "--iamRole",
    dest="iamRole",
    action="store",
    help="IAM role ARN to assume to search for IAM users",
)
parser.add_argument(
    "--allowedUser",
    dest="allowedUsers",
    action="append",
    default=["alpine"],
    help="Allowed users",
)
args = parser.parse_args()

# Fail early if invalid user
if not args.user in args.allowedUsers:
    sys.exit(0)

session = boto3.Session()

if args.iamRole:
    sts = session.client("sts")
    credentials = sts.assume_role(
        RoleArn=args.iamRole, RoleSessionName="sshdKeyLookup"
    )["Credentials"]

    assumed_role_session = boto3.Session(
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )
    iam = assumed_role_session.client("iam")

else:
    iam = session.client("iam")

try:
    for user in iam.get_group(GroupName=args.group)["Users"]:
        for key_desc in iam.list_ssh_public_keys(UserName=user["UserName"])[
            "SSHPublicKeys"
        ]:
            key = iam.get_ssh_public_key(
                UserName=user["UserName"],
                SSHPublicKeyId=key_desc["SSHPublicKeyId"],
                Encoding="SSH",
            )
            if key["SSHPublicKey"]["Status"] == "Active":
                print(key["SSHPublicKey"]["SSHPublicKeyBody"], user["UserName"])
except:
    pass