AWSRecreateSG.py

import demistomock as demisto  # noqa: F401
from CommonServerPython import *  # noqa: F401

from typing import Any
import traceback
from random import randint

ROLE_SESSION_NAME = "xsoar-session"


def split_rule(rule: dict, port: int, protocol: str) -> list[dict]:
    """
    If there are rules with ranges of ports, split them up

    Args:
        rule (Dict): SG rule as a dictionary
        port (int): TCP/UDP Port to be restricted
        protocol (str): Protocol of the port to be restricted

    Returns:
        List: List of replacement rules=
    """
    res_list = []
    # Check if 'FromPort' is in rule, else it is an "all traffic rule".
    if rule.get('FromPort'):
        # Port of interest is in front of back of range, therefore, edit the original rule.
        if rule['FromPort'] == port:
            rule['FromPort'] = rule['FromPort'] + 1
            res_list.append(rule)
        elif rule['ToPort'] == port:
            rule['ToPort'] = rule['ToPort'] - 1
            res_list.append(rule)
        # If in the middle, create two rules.
        else:
            rule_copy = rule.copy()
            rule['ToPort'] = port - 1
            res_list.append(rule)
            rule_copy['FromPort'] = port + 1
            res_list.append(rule_copy)
    else:
        # Splitting up "all traffic" rules.
        if protocol == 'tcp':
            res_list = [{'IpProtocol': 'tcp', 'IpRanges': [{'CidrIp': '0.0.0.0/0'}], 'Ipv6Ranges': [], 'PrefixListIds': [],
                         'UserIdGroupPairs': [], 'FromPort': 0, 'ToPort': port - 1},
                        {'IpProtocol': 'tcp', 'IpRanges': [{'CidrIp': '0.0.0.0/0'}], 'Ipv6Ranges': [], 'PrefixListIds': [],
                         'UserIdGroupPairs': [], 'FromPort': port + 1, 'ToPort': 65535},
                        {'IpProtocol': 'udp', 'IpRanges': [{'CidrIp': '0.0.0.0/0'}], 'Ipv6Ranges': [], 'PrefixListIds': [],
                         'UserIdGroupPairs': [], 'FromPort': 0, 'ToPort': 65535}]
        else:
            res_list = [{'IpProtocol': 'udp', 'IpRanges': [{'CidrIp': '0.0.0.0/0'}], 'Ipv6Ranges': [], 'PrefixListIds': [],
                         'UserIdGroupPairs': [], 'FromPort': 0, 'ToPort': port - 1},
                        {'IpProtocol': 'udp', 'IpRanges': [{'CidrIp': '0.0.0.0/0'}], 'Ipv6Ranges': [], 'PrefixListIds': [],
                         'UserIdGroupPairs': [], 'FromPort': port + 1, 'ToPort': 65535},
                        {'IpProtocol': 'tcp', 'IpRanges': [{'CidrIp': '0.0.0.0/0'}], 'Ipv6Ranges': [], 'PrefixListIds': [],
                         'UserIdGroupPairs': [], 'FromPort': 0, 'ToPort': 65535}]
    return (res_list)


def sg_fix(sg_info: list, port: int, protocol: str, assume_role: str, instance_to_use: str, region: str) -> dict:
    """
    For a SG determine what needs to be recreated.
    Calls split_rule() if there are rules with ranges of ports to be split up

    Args:
        sg_info (List): List of information pulled from "aws-ec2-describe-security-groups" command
        port (int): TCP/UDP Port to be restricted
        protocol (str): Protocol of the port to be restricted
        instance_to_use (str): what integration instance to use

    Returns:
        Dict: Dict of the new SG to be used
    """
    info = dict_safe_get(sg_info, (0, 'Contents', 0))
    recreate_list = []
    # Keep track of change in SG or not.
    change = False
    for rule in info['IpPermissions']:
        if rule.get('IpRanges') and len(rule.get('IpRanges')) > 0:
            # Check if 'FromPort' is in rule, else it is an "all traffic rule".
            if rule.get('FromPort'):
                # Don't recreate if it targets just the port of interest.
                if (
                    rule['FromPort'] == port
                    and port == rule['ToPort']
                    and any(d["CidrIp"] == "0.0.0.0/0" for d in rule["IpRanges"])
                    and rule['IpProtocol'] == protocol
                ):
                    change = True
                elif (
                    rule["FromPort"] == port and port == rule["ToPort"]
                    and any(d["CidrIp"] == "0.0.0.0/0" for d in rule["IpRanges"])
                    and rule["IpProtocol"] == protocol
                ):
                    # If condition to check for Quad 0 in the rules list for matching port.
                    change = True
                elif (
                    rule['FromPort'] <= port and port <= rule['ToPort']
                    and any(d["CidrIp"] == "0.0.0.0/0" for d in rule["IpRanges"])
                    and rule['IpProtocol'] == protocol
                ):  # noqa: E127
                    fixed = split_rule(rule, port, protocol)
                    for rule_fix in fixed:
                        new_rule = (str([rule_fix])).replace("'", "\"")
                        recreate_list.append(new_rule)
                        change = True
                else:
                    new_rule = (str([rule])).replace("'", "\"")
                    recreate_list.append(new_rule)
            elif rule.get('IpRanges') and any(d["CidrIp"] == "0.0.0.0/0" for d in rule["IpRanges"]):
                fixed = split_rule(rule, port, protocol)
                change = True
                for rule_fix in fixed:
                    new_rule = (str([rule_fix])).replace("'", "\"")
                    recreate_list.append(new_rule)
            else:
                new_rule = (str([rule])).replace("'", "\"")
                recreate_list.append(new_rule)
    if change is False:
        return {}
    else:
        # Add rules that allow private_ips to specific port.
        priv_ips_list = [[{'IpProtocol': protocol, 'IpRanges': [{'CidrIp': '10.0.0.0/8'}], 'Ipv6Ranges': [],
                           'PrefixListIds': [], 'UserIdGroupPairs': [], 'FromPort': port, 'ToPort': port}],
                         [{'IpProtocol': protocol, 'IpRanges': [{'CidrIp': '172.16.0.0/12'}], 'Ipv6Ranges': [],
                             'PrefixListIds': [], 'UserIdGroupPairs': [], 'FromPort': port, 'ToPort': port}],
                         [{'IpProtocol': protocol, 'IpRanges': [{'CidrIp': '192.168.0.0/16'}], 'Ipv6Ranges': [],
                           'PrefixListIds': [], 'UserIdGroupPairs': [], 'FromPort': port, 'ToPort': port}]]
        for priv in priv_ips_list:
            recreate_list.append(str(priv).replace("'", "\""))
        new_name = info['GroupName'] + "_xpanse_ar_" + str(randint(100, 999))
        description = "copied from rule " + info['GroupName'] + " by Xpanse Active Response module"
        cmd_args = {"groupName": new_name, "vpcId": info['VpcId'], "description": description, "using": instance_to_use}
        if assume_role:
            cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
        if region:
            cmd_args.update({'region': region})
        new_sg = demisto.executeCommand("aws-ec2-create-security-group", cmd_args)
        if isError(new_sg):
            raise ValueError('Error on creating new security group')
        new_id = dict_safe_get(new_sg, (0, 'Contents', 'GroupId'))
        print(f'{new_id=}')
    for item in recreate_list:
        cmd_args = {"groupId": new_id, "IpPermissionsFull": item, "using": instance_to_use}
        if assume_role:
            cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
        if region:
            cmd_args.update({'region': region})
        res = demisto.executeCommand("aws-ec2-authorize-security-group-ingress-rule",
                                     cmd_args)
        if isError(res):
            if "already exists" in res[0]['Contents']:
                pass
            else:
                raise ValueError('Error on adding security group ingress rules to new security group')
    # Check if there was a rule for `all traffic` (added by default), but break if it is the only egress rule.
    match_all_trafic = False
    for egress in info['IpPermissionsEgress']:
        if egress["IpProtocol"] == '-1':
            if len(info['IpPermissionsEgress']) == 1:
                break
            match_all_trafic = True
        e_format = str([egress]).replace("'", "\"")
        cmd_args = {"groupId": new_id, "IpPermissionsFull": e_format, "using": instance_to_use}
        if assume_role:
            cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
        if region:
            cmd_args.update({'region': region})
        res = demisto.executeCommand("aws-ec2-authorize-security-group-egress-rule",
                                     cmd_args)
        # Don't error if the message is that the rule already exists.
        if isError(res):
            if "already exists" in res[0]['Contents']:
                pass
            else:
                raise ValueError('Error on adding security group egress rules to new security group')
    # If `all traffic` rule before, remove the default one.
    if match_all_trafic is True:
        all_traffic_rule = """[{"IpProtocol": "-1", "IpRanges": [{"CidrIp": "0.0.0.0/0"}], "Ipv6Ranges": [],""" + \
                           """"PrefixListIds": [], "UserIdGroupPairs": []}]"""
        cmd_args = {"groupId": new_id, "IpPermissionsFull": all_traffic_rule, "using": instance_to_use}
        if assume_role:
            cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
        if region:
            cmd_args.update({'region': region})
        res = demisto.executeCommand("aws-ec2-revoke-security-group-egress-rule",
                                     cmd_args)
        if isError(res):
            raise ValueError('Error on removing egress `allow all` rule on new security group')
    return {'new-sg': new_id}


def replace_sgs(replace_list: list, int_sg_mapping: dict, assume_role: str, instance_to_use: str, region: str):
    """
    Replace the actual SGs on the interface

    Args:
        replace_list (List): list of SGs to be replaced
        int_sg_mapping: interface to security group mapping
        instance_to_use (str): what integration instance to use

    Returns:
        none
    """
    for entry in replace_list:
        int_sg_mapping[entry['int']].remove(entry['old-sg'])
        int_sg_mapping[entry['int']].append(entry['new-sg'])
        formatted_list = ','.join(int_sg_mapping[entry['int']])
        cmd_args = {"networkInterfaceId": entry['int'], "groups": formatted_list, "using": instance_to_use}
        if assume_role:
            cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
        if region:
            cmd_args.update({'region': region})
        res = demisto.executeCommand("aws-ec2-modify-network-interface-attribute",
                                     cmd_args)
        if isError(res):
            raise ValueError('Error on replacing security group(s) on network interface')


def determine_excessive_access(int_sg_mapping: dict, port: int, protocol: str, assume_role: str, instance_to_use: str,
                               region: str) -> list:
    """
    Pulls info on each SG and then calls sg_fix() to actually create the new SGs

    Args:
        int_sg_mapping (Dict): interface to security group mapping
        port (int): TCP/UDP Port to be restricted
        protocol (str): Protocol of the port to be restricted
        instance_to_use (str): what integration instance to use

    Returns:
        List: list of SGs to be replaced
    """
    replace_list = []
    for mapping in int_sg_mapping:
        for sg in int_sg_mapping[mapping]:
            cmd_args = {"groupIds": sg, "using": instance_to_use}
            if region:
                cmd_args.update({'region': region})
            if assume_role:
                cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
            sg_info = demisto.executeCommand("aws-ec2-describe-security-groups", cmd_args)
            if isError(sg_info):
                raise ValueError('Error on describing security group')
            elif sg_info:
                res = sg_fix(sg_info, port, protocol, assume_role, instance_to_use, region)
                # Need interface, old sg and new sg.
                if res.get('new-sg'):
                    res['old-sg'] = sg
                    res['int'] = mapping
                    replace_list.append(res)
    return replace_list


def instance_info(instance_id: str, public_ip: str, assume_role: str, region: str) -> tuple[dict, str]:
    """
    Finds interface with public_ip and from this creates interface ID/SG mapping

    Args:
        instance_id (str): EC2 Instance ID
        public_ip (str): Public IP address of the EC2 instance

    Returns:
        tuple[dict, str]: A dictionary mapping interfaces to security groups (dict), and an integration to use (str).
    """
    cmd_args = {"instanceIds": instance_id}
    if region:
        cmd_args.update({'region': region})
    if assume_role:
        cmd_args.update({'roleArn': assume_role, 'roleSessionName': ROLE_SESSION_NAME})
    instance_info = demisto.executeCommand("aws-ec2-describe-instances", cmd_args)
    # Need a for loop in case multiple AWS-EC2 integrations are configured.
    match = False
    for instance in instance_info:
        interfaces = dict_safe_get(instance, ('Contents', 0, 'NetworkInterfaces'))
        if not isError(instance) and interfaces:
            mapping_dict = {}
            for interface in interfaces:
                if interface.get('Association') and interface.get('Association').get('PublicIp') == public_ip:
                    match = True
                    group_list = []
                    for sg in interface['Groups']:
                        group_list.append(sg['GroupId'])
                    mapping_dict[interface['NetworkInterfaceId']] = group_list
                    instance_to_use = instance['Metadata']['instance']
                    break
        if match:
            break
    if match is False:
        raise ValueError('could not find interface with public IP association')
    return mapping_dict, instance_to_use


def create_command_results(readable_output: str, output_flag: bool):
    command_results = CommandResults(
        outputs={'awssgrecreated': output_flag},
        raw_response={'awssgrecreated': output_flag},
        readable_output=readable_output,
    )
    return command_results


def aws_recreate_sg(args: dict[str, Any]) -> str:
    """
    Main command that determines what interface on an EC2 instance has an over-permissive security group on,
    determine which security groups have over-permissive rules and to replace them with a copy of the security group
    that has only the over-permissive portion removed.  Over-permissive is defined as sensitive ports (SSH, RDP, etc)
    being exposed to the internet via IPv4.

    Args:
        args (Dict[str, Any]): Demisto.args() object

    Returns:
        str: human readable message of what SGs were replaced on what interface
    """

    instance_id = args.get('instance_id', None)
    port = int(args.get('port', None))
    protocol = args.get('protocol', None)
    public_ip = args.get('public_ip', None)
    assume_role = args.get('assume_role', None)
    region = args.get('region', None)

    if not instance_id or not port or not protocol or not public_ip:
        raise ValueError('instance_id, port, protocol and public_ip all need to be specified')

    # Determine interface with public IP and associated SGs.
    int_sg_mapping, instance_to_use = instance_info(instance_id, public_ip, assume_role, region)
    # Determine what SGs are overpermissive for particular port.
    replace_list = determine_excessive_access(int_sg_mapping, port, protocol, assume_role, instance_to_use, region)
    if len(replace_list) == 0:
        readable_output = 'No security groups were found to need to be replaced'
        return create_command_results(readable_output, False)
    replace_sgs(replace_list, int_sg_mapping, assume_role, instance_to_use, region)
    readable_output = f"For interface {replace_list[0]['int']}: \r\n"
    for replace in replace_list:
        readable_output += f"replaced SG {replace['old-sg']} with {replace['new-sg']} \r\n"
    return create_command_results(readable_output, True)


''' MAIN FUNCTION '''


def main():
    try:
        return_results(aws_recreate_sg(demisto.args()))
    except Exception as ex:
        demisto.error(traceback.format_exc())  # print the traceback
        return_error(f'Failed to execute AWSRecreateSG. Error: {str(ex)}')


''' ENTRY POINT '''


if __name__ in ('__main__', '__builtin__', 'builtins'):
    main()