#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This CLI tool can be used to scan for dangling elastic IPs in your AWS accounts by correlating data
betwen Route53 records and Elastic IPs / IPs in your AWS account.

.. currentmodule:: ghostbusterV4.scan
.. moduleauthor:: shubham_shah <sshah@assetnote.io>
"""
import csv
import datetime
import sys
import requests
import click
from .__init__ import Info, pass_info
import boto3
import base64
import json as json_lib
import CloudFlare
import awsipranges
from slack_sdk.webhook import WebhookClient
from botocore.exceptions import ClientError

route53_zone_records_set = {}

@click.group(
    help="Commands that help you scan your AWS account for dangling elastic IPs"
)
@click.pass_context
def cli(ctx: click.Context):
    """CLI handler for scanning actions"""
    pass

# ascii art

logo_b64 = " "
json_output = False


def log(text):
    if not json_output:
        click.echo(text)

def send_webhook(slackwebhook, takeovers):
    webhook = WebhookClient(url=slackwebhook)
    from slack_sdk.http_retry.builtin_handlers import RateLimitErrorRetryHandler

    rate_limit_handler = RateLimitErrorRetryHandler(max_retry_count=1)
    webhook.retry_handlers.append(rate_limit_handler)
    for takeover in takeovers:
        payload = "Potential AWS Elastic IP takeover: {} Records: {}".format(
            takeover["name"], takeover["records"]
        )
        _ = webhook.send(text=payload)

def get_cloudflare_records(cloudflaretoken):
    log("Obtaining all zone names from Cloudflare.")
    cf = CloudFlare.CloudFlare(token=cloudflaretoken, raw=True)
    dns_records = []
    # get zone names
    cloudflare_zones = []
    try:
        page_number = 0
        while True:
            page_number += 1
            raw_results = cf.zones.get(
                params={"per_page": 100, "page": page_number}
            )
            zones = raw_results["result"]
            for zone in zones:
                zone_id = zone["id"]
                zone_name = zone["name"]
                cloudflare_zones.append({"name": zone_name, "id": zone_id})
            total_pages = raw_results["result_info"]["total_pages"]
            if page_number == total_pages:
                break
    except CloudFlare.exceptions.CloudFlareAPIError as e:
        exit("Failed to retreive zones %d %s - api call failed" % (e, e))

    log("Obtaining DNS A records for all zones from Cloudflare.")
    # get dns records for zones
    for zone in cloudflare_zones:
        try:
            page_number = 0
            while True:
                page_number += 1
                raw_results = cf.zones.dns_records.get(
                    zone["id"], params={"per_page": 100, "page": page_number}
                )
                cf_dns_records = raw_results["result"]
                for record in cf_dns_records:
                    if record.get("content"):
                        if record["type"] == "A":
                            dns_records.append(
                                {
                                    "name": record["name"],
                                    "records": [record["content"]],
                                }
                            )
                total_pages = raw_results["result_info"]["total_pages"]
                if page_number == total_pages:
                    break
        except CloudFlare.exceptions.CloudFlareAPIError as e:
            exit("Failed to retreive DNS records %d %s - api call failed" % (e, e))
    return dns_records

def get_route53_hosted_zones(route53, next_zone=None):
    """Recursively returns a list of hosted zones in Amazon Route 53."""
    if next_zone:
        response = route53.list_hosted_zones_by_name(
            DNSName=next_zone[0], HostedZoneId=next_zone[1]
        )
    else:
        response = route53.list_hosted_zones_by_name()
    hosted_zones = response["HostedZones"]
    # if response is truncated, call function again with next zone name/id
    if response["IsTruncated"]:
        hosted_zones += get_route53_hosted_zones(
            route53, (response["NextDNSName"], response["NextHostedZoneId"])
        )
    return hosted_zones

def get_route53_zone_records(route53, zone_id, next_record=None):
    """Recursively returns a list of records of a hosted zone in Route 53."""
    if next_record:
        response = route53.list_resource_record_sets(
            HostedZoneId=zone_id,
            StartRecordName=next_record[0],
            StartRecordType=next_record[1],
        )
    else:
        response = route53.list_resource_record_sets(HostedZoneId=zone_id)
    zone_records = response["ResourceRecordSets"]

    # if response is truncated, call function again with next record name/id
    if response["IsTruncated"]:
        zone_set_record_key = response["NextRecordName"] + "_" + response["NextRecordType"]
        if zone_set_record_key in route53_zone_records_set == False:
            zone_records += get_route53_zone_records(
                route53, zone_id, (response["NextRecordName"], response["NextRecordType"])
            )
            route53_zone_records_set.add(zone_set_record_key)
    return zone_records

def get_record_value(record):
    """Return a list of values for a hosted zone record."""
    # test if record's value is Alias or dict of records
    try:
        value = [
            ":".join(
                [
                    "ALIAS",
                    record["AliasTarget"]["HostedZoneId"],
                    record["AliasTarget"]["DNSName"],
                ]
            )
        ]
    except KeyError:
        value = []
        for v in record["ResourceRecords"]:
            value.append(v["Value"])
    return value

def try_record(test, record):
    """Return a value for a record"""
    # test for Key and Type errors
    try:
        value = record[test]
    except KeyError:
        value = ""
    except TypeError:
        value = ""
    return value

def get_route53_A_records(route53):
    dns_records = []
    hosted_zones = get_route53_hosted_zones(route53)
    for zone in hosted_zones:
        zone_records = get_route53_zone_records(route53, zone["Id"])
        for record in zone_records:
            if record["Type"] == "A":
                # we aren't interested in alias records
                if record.get("AliasTarget"):
                    # skip
                    pass
                else:
                    a_records = []
                    for r53value in record["ResourceRecords"]:
                        a_records.append(r53value["Value"])
                    r53_obj = {"name": record["Name"], "records": a_records}
                    dns_records.append(r53_obj)
    return dns_records

def get_eips(ec2, region):
    # super annoying, boto3 doesn't have a native paginator class for describe_addresses
    elastic_ips = []
    while True:
        addresses_dict = []
        if addresses_dict and "NextToken" in addresses_dict:
            addresses_dict = client.describe_addresses(
                NextToken=addresses_dict["NextToken"]
            )
        else:
            addresses_dict = ec2.describe_addresses()
        for eip_dict in addresses_dict["Addresses"]:
            elastic_ips.append(eip_dict["PublicIp"])
        if "NextToken" not in addresses_dict:
            break

    log("Obtaining IPs for network interfaces for region: {}".format(region))
    nic_paginator = ec2.get_paginator("describe_network_interfaces")
    for resp in nic_paginator.paginate():
        for interface in resp.get("NetworkInterfaces", []):
            if interface.get("Association"):
                nic_public_ip = interface["Association"]["PublicIp"]
                elastic_ips.append(nic_public_ip)
    return elastic_ips

def assume_role(role_arn):
    session = boto3.Session()
    sts = session.client("sts")
    assumed_role_object = sts.assume_role(
        RoleArn=role_arn,
        RoleSessionName="ghostbusterV4-session"
    )
    assumed_role_credentials = assumed_role_object["Credentials"]

    return boto3.Session(
        aws_access_key_id=assumed_role_credentials["AccessKeyId"],
        aws_secret_access_key=assumed_role_credentials["SecretAccessKey"],
        aws_session_token=assumed_role_credentials["SessionToken"]
    )

def get_all_account_ids(organisation_lookup_role_arn):
    account_ids = []
    session = assume_role(organisation_lookup_role_arn)
    organizations = session.client("organizations")

    def add_account_ids(list_accounts_response):
        for account in list_accounts_response["Accounts"]:
            account_ids.append(account["Id"])

    response = organizations.list_accounts()
    add_account_ids(response)
    while "NextToken" in response:
        response = organizations.list_accounts(NextToken=response["NextToken"])
        add_account_ids(response)
    return account_ids

def check_whitelist(whitelist_ips,record):
    if record in whitelist_ips:
        return True
    else:
        return False

def parseWhitelistFile(whitelist):
    whitelist_ips = []
    if whitelist != "" :
        with open(whitelist,"r") as f:
            contents = f.readlines()
            for ip in contents:
                ip.replace(" ","")
                ip=ip[:-1]
                whitelist_ips.append(ip)
    return whitelist_ips

@click.option(
    "--regions", default="us-east-1", help="Comma delimited list of regions to run on."
)
@click.option(
    "--exclude", default="", help="Comma delimited list of profile names to exclude."
)
@click.option(
    "--allregions",
    default=False,
    is_flag=True,
    help="Run on all regions.",
)
@click.option(
    "--cloudflaretoken",
    default="",
    help="Pull DNS records from Cloudflare, provide a CF API token.",
)
@click.option(
    "--records",
    required=False,
    type=click.Path(exists=True),
    help="Manually specify DNS records to check against. ghostbusterV4 will check these IPs after checking retrieved DNS "
         "records. See records.csv for an example.",
)
@click.option(
    "--slackwebhook",
    default="",
    help="Specify a Slack webhook URL to send notifications about potential takeovers. Slack workflow variable must be"
         "named 'text'",
)
@click.option(
    "--json",
    default=False,
    is_flag=True,
    help="Only return a JSON object.",
)
@click.option(
    "--skipascii",
    default=False,
    is_flag=True,
    help="Skip printing the ASCII art when starting up ghostbusterV4.",
)
@click.option(
    "--profile",
    default="",
    help="Specify a specific AWS profile to run ghostbusterV4 on.",
)
@click.option(
    "--roles",
    required=False,
    type=click.Path(exists=True),
    help="Specify CSV filename with AWS account IDs to run ghostbusterV4 on. Each account must have ghostbusterV4 role "
         "assumable by ghostbusterV4 ec2/lambda/whatever is running ghostbusterV4. Role name: ghostbusterV4TargetAccountRole."
         " See roles.csv for example.",
)
@click.option(
    "--autoroles",
    default="",
    required=False,
    help="Like --roles, but finds all organisation accounts automatically. The argument value should be ARN of a role "
         "with organizations:ListAccounts and organizations:DescribeAccount. Ec2/lambda/whatever is running ghostbusterV4"
         " must have permissions to assume the organisation lookup role."
)
@click.option(
    "--whitelist", default="", help="Specify filepath having list of IPs to whitelist"
)
@cli.command(help="Scan for dangling elastic IPs inside your AWS accounts.")
@pass_info
def aws(
    _: Info,
    regions: str,
    exclude: str,
    allregions: bool,
    cloudflaretoken: str,
    records: str,
    slackwebhook: str,
    skipascii: str,
    profile: str,
    roles: str,
    autoroles: str,
    json: bool,
    whitelist: str
    ):
    """Scan for dangling elastic IPs inside your AWS accounts."""
    # ascii art
    if not skipascii and not json:
        sys.stdout.write(base64.b64decode(logo_b64).decode('utf-8'))
    
    global json_output
    json_output = json
        
    session = boto3.Session()
    profiles = session.available_profiles

    if exclude != "":
        exclude_list = exclude.split(",")
        for excluded_profile in exclude_list:
            profiles.remove(excluded_profile)
    if profile != "":
        profiles = [profile]

    account_ids = []
    if roles:
        account_ids = [account_id["account_id"] for account_id in csv.DictReader(open(roles, "r"))]
    elif autoroles:
        log("Finding accounts automatically using role: {0}".format(autoroles))
        account_ids = get_all_account_ids(autoroles)
        log("Found {0} accounts in the organisation.".format(len(account_ids)))

    dns_records = []
    # collection of records from cloudflare
    if cloudflaretoken != "":
        cf_dns_records = get_cloudflare_records(cloudflaretoken)
        dns_records = dns_records + cf_dns_records
        log("Obtained {0} DNS A records so far.".format(len(dns_records)))

    # collection of records from r53 using profiles
    for profile in profiles:
        profile_session = boto3.session.Session(profile_name=profile)
        route53 = profile_session.client("route53")
        log("Obtaining Route53 hosted zones for AWS profile: {0}.".format(profile))
        dns_records.extend(get_route53_A_records(route53))

    # collection of records from r53 using roles
    for account_id in account_ids:
        role_arn = "arn:aws:iam::{0}:role/ghostbusterV4TargetAccountRole".format(account_id)
        try:
            role_session = assume_role(role_arn)
            log("Successfully assumed role from account {0}".format(account_id))
        except ClientError as error:
            log("Failed to assume role {0}, skipping it. Error: {1}".format(role_arn, error))
            continue
        route53 = role_session.client('route53')
        log("Obtaining Route53 hosted zones for AWS account ID: {0}.".format(account_id))
        dns_records.extend(get_route53_A_records(route53))

    log("Obtained {0} DNS A records so far.".format(len(dns_records)))

    # collection of IPs
    if allregions:
        ec2 = boto3.client("ec2")
        aws_regions = [
            region["RegionName"] for region in ec2.describe_regions()["Regions"]
        ]
    else:
        aws_regions = regions.split(",")

    whitelist_ips = parseWhitelistFile(whitelist)
    elastic_ips = []
    # collect elastic compute addresses / EIPs for all regions
    for region in aws_regions:
        for profile in profiles:
            log("Obtaining EIPs for region: {}, profile: {}".format(region, profile))
            profile_session = boto3.session.Session(profile_name=profile)
            ec2 = profile_session.client("ec2", region_name=region)
            elastic_ips.extend(get_eips(ec2=ec2, region=region))
        for account_id in account_ids:
            log("Obtaining EIPs for region: {}, account ID: {}".format(region, account_id))
            role_arn = "arn:aws:iam::{0}:role/ghostbusterV4TargetAccountRole".format(account_id)
            try:
                role_session = assume_role(role_arn)
            except ClientError as error:
                log("Failed to assume role {0}, skipping it. Error: {1}".format(role_arn, error))
                continue

            ec2 = role_session.client("ec2", region_name=region)
            elastic_ips.extend(get_eips(ec2=ec2, region=region))

    unique_ips = list(set(elastic_ips))
    log("Obtained {0} unique elastic IPs from AWS.".format(len(unique_ips)))

    dns_ec2_ips = []
    # find all DNS records that point to EC2 IP addresses
    aws_ip_ranges = awsipranges.get_ranges()
    for record_set in dns_records:
        for record in record_set["records"]:
            aws_metadata = aws_ip_ranges.get(record)
            if aws_metadata:
                for service in aws_metadata.services:
                    if service == "EC2":
                        dns_ec2_ips.append(record_set)

    takeovers = []
    # check to see if any of the record sets we have, we don't own the elastic IPs
    for record_set in dns_ec2_ips:
        for record in record_set["records"]:
            if(not check_whitelist(whitelist_ips,record)):
                if record not in elastic_ips:
                    takeovers.append(record_set)
                    log("Takeover possible: {}".format(record_set))

    # check if manually specified A records exist in AWS acc (eips/public ips)
    if records:
        with open(records, "r") as fp:
            csv_reader = csv.DictReader(fp)
            for row in csv_reader:
                aws_metadata = aws_ip_ranges.get(row["record"])
                if aws_metadata:
                    for service in aws_metadata.services:
                        if service == "EC2":
                            if(not check_whitelist(whitelist_ips,row["record"])):
                                if row["record"] not in elastic_ips:
                                    takeover_obj = {
                                        "name": row["name"],
                                        "records": [row["record"]],
                                    }
                                    takeovers.append(takeover_obj)
                                    log("Takeover possible: {}".format(takeover_obj))

    if len(takeovers) == 0:
        log("No takeovers detected! Nice work.")
    if json:
        click.echo(json_lib.dumps(takeovers, indent=2))

    # send slack webhooks, with retries in case of 429s
    if slackwebhook != "":
        send_webhook(slackwebhook, takeovers)
