import demistomock as demisto
from CommonServerPython import *  # noqa # pylint: disable=unused-wildcard-import
from CommonServerUserPython import *  # noqa
from AWSApiModule import *  # noqa: E402
from typing import Callable, Tuple
import urllib3.util
import boto3
import urllib3

# Disable insecure warnings
urllib3.disable_warnings()

''' CONSTANTS '''

SERVICE = 'wafv2'
OUTPUT_PREFIX = 'AWS.Waf'
DEFAULT_SCOPE = 'Regional'
REGEX_MATCH_STATEMENT = 'RegexPatternSetReferenceStatement'
BYTE_MATCH_STATEMENT = 'ByteMatchStatement'
TEXT_TRANSFORMATIONS = 'NONE | COMPRESS_WHITE_SPACE | HTML_ENTITY_DECODE | LOWERCASE | CMD_LINE | URL_DECODE | ' \
                       'BASE64_DECODE | HEX_DECODE | MD5 | REPLACE_COMMENTS | ESCAPE_SEQ_DECODE | SQL_HEX_DECODE | ' \
                       'CSS_DECODE | JS_DECODE | NORMALIZE_PATH | NORMALIZE_PATH_WIN | REMOVE_NULLS | ' \
                       'REPLACE_NULLS | BASE64_DECODE_EXT | URL_DECODE_UNI | UTF8_TO_UNICODE'

SCOPE_MAP = {'Regional': 'REGIONAL',
             'Global': 'CLOUDFRONT'}
OPERATOR_TO_STATEMENT_OPERATOR = {'And': 'AndStatement', 'Or': 'OrStatement', 'Not': 'NotStatement'}
MATCH_TYPE_TO_POSITIONAL_CONSTRAIN = {'Exactly Matches String': 'EXACTLY',
                                      'Starts With String': 'STARTS_WITH',
                                      'Ends With String': 'ENDS_WITH',
                                      'Contains String': 'CONTAINS',
                                      'Contains Words': 'CONTAINS_WORD',
                                      'all': 'EXACTLY | STARTS_WITH | ENDS_WITH | CONTAINS | CONTAINS_WORD'}
WEB_REQUEST_COMPONENT_MAP = {"Headers": "Headers",
                             "Cookies": "Cookies",
                             "Query Parameters": "AllQueryArguments",
                             "Uri Path": "UriPath",
                             "Query String": "QueryString",
                             "Body": "Body",
                             "HTTP Method": "Method"}

''' HELPER FUNCTIONS '''


def build_string_match_rule_object(args: dict) -> dict:  # pragma: no cover
    """
    Creates a string match rule statement object that can be added to a rule
    Args:
        args: The command arguments

    Returns:
        String match rule statement object
    """
    return {
        'Statement': build_string_match_statement(**args)
    }


def get_required_args_for_get_rule_group(args: dict) -> dict:  # pragma: no cover
    """
    Build the required arguments for a request of get_rule_group
    Args:
        args: The command arguments

    Returns:
       The required arguments for a request of get_rule_group
    """
    return {
        'Name': args.get('group_name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('group_id', '')
    }


def build_ip_statement(ip_set_arn: str) -> dict:  # pragma: no cover
    """
    Creates an ip statement that can be added to a statements list of a rule
    Args:
        ip_set_arn: The ip set ARN representation

    Returns:
        An ip statement object
    """
    return {'IPSetReferenceStatement': {'ARN': ip_set_arn}}


def build_country_statement(country_codes: list) -> dict:  # pragma: no cover
    """
    Creates a country statement that can be added to a statements list of a rule
    Args:
        country_codes: The country codes

    Returns:
        A country statement object
    """
    return {'GeoMatchStatement': {'CountryCodes': country_codes}}


def build_country_rule_object(args: dict) -> dict:  # pragma: no cover
    """
    Creates a country rule statement object that can be added to a rule
    Args:
        args: The command arguments

    Returns:
        Country rule statement object
    """
    country_codes = argToList(args.get('country_codes'))
    return {
        'Statement': build_country_statement(country_codes)
    }


def build_visibility_config_object(metric_name: str,
                                   cloud_watch_metrics_enabled: bool,
                                   sampled_requests_enabled: bool) -> dict:  # pragma: no cover
    """
    Creates a dictionary which represents visibility config
    Args:
        metric_name: The metric name
        cloud_watch_metrics_enabled: whether to enable cloud metrics
        sampled_requests_enabled: whether to enable sample requests

    Returns:
        Visibility config object
    """
    return {
        'CloudWatchMetricsEnabled': cloud_watch_metrics_enabled,
        'MetricName': metric_name,
        'SampledRequestsEnabled': sampled_requests_enabled
    }


def get_tags_dict_from_args(tag_keys: list, tag_values: list) -> list:
    """
    Creates a list of dictionaries containing the tag key, and its corresponding value
    Args:
        tag_keys: tag keys list
        tag_values: tag values list

    Returns:
        List of tags
    """
    if len(tag_keys) != len(tag_values):
        raise DemistoException('The tags_keys and tag_values arguments must be at the same length.')

    # keys and values are in the same length
    return [{'Key': k, 'Value': v} for k, v in zip(tag_keys, tag_values)]


def build_regex_pattern_object(regex_patterns: list) -> List[dict]:
    """
    Creates a list of dictionaries which represent a regex set object
    Args:
        regex_patterns: regex patterns

    Returns:
        List of regex patterns objects
    """
    return [
        {'RegexString': regex_pattern} for regex_pattern in regex_patterns
    ]


def build_ip_rule_object(args: dict) -> dict:
    """
    Creates an ip rule statement object that can be added to a rule
    Args:
        args: The command arguments

    Returns:
        Ip rule statement object
    """
    ip_rule = {}
    ip_set_arns = argToList(args.get('ip_set_arn'))
    condition_operator = args.get('condition_operator', '')
    if len(ip_set_arns) > 1 and not condition_operator:
        raise DemistoException('The condition_operator argument must be specified when '
                               'ip_set_arn contains more than one value.')

    if len(ip_set_arns) == 1:
        ip_rule['Statement'] = build_ip_statement(ip_set_arns[0])
    elif len(ip_set_arns) > 1:
        statement_operator = OPERATOR_TO_STATEMENT_OPERATOR[condition_operator]
        ip_rule.setdefault('Statement', {})[statement_operator] = {
            'Statements': [build_ip_statement(ip_set_arn) for ip_set_arn in ip_set_arns]
        }
    return ip_rule


def build_string_match_statement(match_type: str = '',
                                 string_to_match: str = '',
                                 regex_set_arn: str = '',
                                 oversize_handling: str = '',
                                 text_transformation: str = 'NONE',
                                 web_request_component: str = '', **kwargs) -> dict:
    """
    Creates a byte/regex match statement that can be added to a statements list of a rule
    Args:
        match_type: Which match type should be performed
        string_to_match: string_to_match: The string to match to
        regex_set_arn: The regex set to match to
        oversize_handling: The oversize handling to be applied to web request contents
        text_transformation: The text transformation to perform on the component
        web_request_component: web_request_component: The web request component to inspect
        kwargs: Args that are not in use in this method but passed from the command arguments

    Returns:
        A byte/regex match statement object
    """
    match_statement = REGEX_MATCH_STATEMENT if match_type == 'Matches Regex Pattern Set' \
        else BYTE_MATCH_STATEMENT
    web_request_component = WEB_REQUEST_COMPONENT_MAP.get(web_request_component) or ''
    if match_statement == BYTE_MATCH_STATEMENT:
        if not string_to_match:
            raise DemistoException('string_to_match must be provided when using strings match_type')
        statement = build_byte_match_statement(web_request_component=web_request_component,
                                               oversize_handling=oversize_handling,
                                               text_transformation=text_transformation,
                                               string_to_match=string_to_match,
                                               match_type=match_type)

    else:  # match_statement == REGEX_MATCH_STATEMENT
        if not regex_set_arn:
            raise DemistoException('regex_set_arn must be provided when using Matches Regex Pattern Set match_type')
        statement = build_regex_match_statement(web_request_component=web_request_component,
                                                oversize_handling=oversize_handling,
                                                regex_set_arn=regex_set_arn,
                                                text_transformation=text_transformation)

    return {match_statement: statement}


def update_rule_with_statement(rule: dict, statements: list, condition_operator: str):
    """
    Updates an existing rule with a new statement
    Args:
        rule: The rule to update
        statements: The statement to update the rule with
        condition_operator: The condition to apply on the statements
    """
    old_rule_statement = rule.get('Statement', {})
    if 'AndStatement' in old_rule_statement or 'OrStatement' in old_rule_statement:
        demisto.info('ignoring condition_operator argument as the statement already contains an operator.')
        condition = list(old_rule_statement.keys())[0]
    elif condition_operator:
        condition = OPERATOR_TO_STATEMENT_OPERATOR[condition_operator]
        # override the statement key with the conditional statement
        rule['Statement'] = {condition: {'Statements': [old_rule_statement]}}

    else:
        raise DemistoException('Rule contains only one statement. Please provide condition operator.')
    rule['Statement'][condition]['Statements'].extend(statements)


def create_rules_list_with_new_rule_statement(args: dict, statements: list, rules: list) -> list:
    """
    Creates a rules list with the updated rule
    Args:
        args: The command arguments
        statements: The statements to add to a rule
        rules: The original rules

    Returns:
        Updated list of rules
    """
    new_rules = rules.copy()
    rule_name = args.get('rule_name', '')
    condition_operator = args.get('condition_operator', '')
    for rule in new_rules:
        if rule.get('Name') == rule_name:
            update_rule_with_statement(rule, statements, condition_operator)
            return rules
    raise DemistoException(f'Did not find any rule with name {rule_name}')


def build_web_component_match_object(web_request_component: str, oversize_handling: str) -> dict:
    """
    Creates web component object to send to the API
    Args:
        web_request_component: The web request component to inspect
        oversize_handling: The oversize handling to be applied to web request contents

    Returns:

    """
    web_request_component_object = {}
    if web_request_component in {'Headers', 'Cookies', 'Body'}:
        if not oversize_handling:
            raise DemistoException(
                'oversize_handling must be provided when using Headers, Cookies, Body in web_request_component')
        if web_request_component != 'Body':
            web_request_component_object = {'MatchPattern': {
                'All': {}
            },
                'MatchScope': 'ALL'
            }
        web_request_component_object['OversizeHandling'] = oversize_handling
    return web_request_component_object


def build_byte_match_statement(web_request_component: str,
                               oversize_handling: str,
                               text_transformation: str,
                               string_to_match: str,
                               match_type: str) -> dict:
    """
    Creates a byte match statement
    Args:
        web_request_component: The web request component to inspect
        oversize_handling: The oversize handling to be applied to web request contents
        that are bigger than what can be inspected by AWS WAF
        text_transformation: The text transformation to perform on the component
        string_to_match: The string to match to
        match_type: Which match type should be performed

    Returns:
        A byte match statement value
    """
    web_request_component_object = build_web_component_match_object(web_request_component, oversize_handling)
    return {
        'SearchString': string_to_match,
        'FieldToMatch': {
            web_request_component: web_request_component_object
        },
        'TextTransformations': [
            {'Priority': 0,
             'Type': text_transformation
             }
        ],
        'PositionalConstraint': MATCH_TYPE_TO_POSITIONAL_CONSTRAIN[match_type]}


def build_regex_match_statement(web_request_component: str,
                                oversize_handling: str,
                                text_transformation: str,
                                regex_set_arn: str) -> dict:
    """
    Creates a byte match statement
    Args:
        web_request_component: The web request component to inspect
        oversize_handling: The oversize handling to be applied to web request contents
        that are bigger than what can be inspected by AWS WAF
        text_transformation: The text transformation to perform on the component
        regex_set_arn: The regex set to match to

    Returns:
        A regex match statement value
    """
    web_request_component_object = build_web_component_match_object(web_request_component, oversize_handling)
    return {
        'ARN': regex_set_arn,
        'FieldToMatch': {
            web_request_component: web_request_component_object
        },
        'TextTransformations': [
            {'Priority': 0, 'Type': text_transformation
             }
        ]
    }


def build_new_rule_object(args: dict, rule_group_visibility_config: dict,
                          build_rule_func: Callable[[dict], dict]) -> dict:
    """
    Creates a country rule object that can be added to a rule group rules list
    Args:
        args: The command arguments
        rule_group_visibility_config: The rule visibility config
        build_rule_func: A generic function that builds the statement of the rule

    Returns:
        Entire rule object
    """
    name = args.get('rule_name', '')
    rule_visibility_config = build_visibility_config_object(
        metric_name=name,
        cloud_watch_metrics_enabled=rule_group_visibility_config.get('CloudWatchMetricsEnabled', True),
        sampled_requests_enabled=rule_group_visibility_config.get('SampledRequestsEnabled', True))

    rule = {
        'Name': name,
        'Priority': arg_to_number(args.get('priority', '')) or 0,
        'Action': {
            args.get('action'): {}
        },
        'VisibilityConfig': rule_visibility_config,

    }
    rule |= build_rule_func(args)

    return rule


def delete_rule(rule_name: str, rules: list) -> list:
    """
    Deletes a rule from a rules list
    Args:
        rule_name: The rule name to delete
        rules: The rules list of a rule group

    Returns:
        A new rules list without the rule to delete
    """
    updated_rules = rules.copy()
    for rule in rules:
        if rule.get('Name') == rule_name:
            updated_rules.remove(rule)
            break
    return updated_rules


def append_new_rule(rules: list, rule: dict) -> list:
    """
    Adds a rule from a rules list
    Args:
        rule: The rule object to add
        rules: The rules list of a rule group

    Returns:
        A new rules list with the rule to delete
    """
    updated_rules = rules.copy()
    updated_rules.append(rule)
    return updated_rules


def get_required_response_fields_from_rule_group(client: boto3.client, kwargs: dict) -> Tuple[list, dict, str]:
    """
    Gets all the fields from the response that are required for the update request
    Args:
        client: AWS WF client
        kwargs: args required for get rule group

    Returns:
        rules, visibility config object and lockToken associated to a rule group
    """
    response = client.get_rule_group(**kwargs)

    rule_group = response.get('RuleGroup', {})
    rules = rule_group.get('Rules', [])
    rule_group_visibility_config = rule_group.get('VisibilityConfig', {})
    lock_token = response.get('LockToken')

    return rules, rule_group_visibility_config, lock_token


'''CLIENT FUNCTIONS'''


def update_rule_group_rules(client: boto3.client,
                            kwargs: dict,
                            lock_token: str,
                            updated_rules: list,
                            rule_group_visibility_config: dict) -> dict:  # pragma: no cover
    """ Updates rule group with new rules list"""
    kwargs |= {'LockToken': lock_token,
               'Rules': updated_rules,
               'VisibilityConfig': rule_group_visibility_config
               }

    return client.update_rule_group(**kwargs)


''' COMMAND FUNCTIONS '''


def connection_test(client: boto3.client) -> str:  # pragma: no cover
    """ Command to test the connection to the API"""
    try:
        client.list_ip_sets(Scope=SCOPE_MAP[DEFAULT_SCOPE])
    except Exception as e:
        raise DemistoException(f'Failed to execute test module. Error: {str(e)}')

    return 'ok'


def create_ip_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to create an IP set"""
    tag_keys = argToList(args.get('tag_key')) or []
    tag_values = argToList(args.get('tag_value')) or []
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'IPAddressVersion': args.get('ip_version', ''),
        'Addresses': argToList(args.get('addresses')) or [],
    }

    if description := args.get('description'):
        kwargs |= {'Description': description}
    if tags := get_tags_dict_from_args(tag_keys, tag_values):
        kwargs |= {'Tags': tags}

    response = client.create_ip_set(**kwargs)
    outputs = response.get('Summary', {})

    readable_output = f'AWS Waf ip set with id {outputs.get("Id", "")} was created successfully'

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_prefix=f'{OUTPUT_PREFIX}.IpSet',
                          outputs_key_field='Id')


def get_ip_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to get a specific IP set"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    response = client.get_ip_set(**kwargs)

    outputs = response.get('IPSet', {})

    readable_output = tableToMarkdown('IP Set', outputs)

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_prefix=f'{OUTPUT_PREFIX}.IpSet',
                          outputs_key_field='Id')


def update_ip_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to update a specific IP set"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    addresses_to_update = argToList(args.get('addresses')) or []
    overwrite = argToBoolean(args.get('is_overwrite')) or False

    get_response = client.get_ip_set(**kwargs)

    lock_token = get_response.get('LockToken', '')
    original_addresses = get_response.get('IPSet', {}).get('Addresses')
    if not overwrite:
        addresses_to_update.extend(original_addresses)

    kwargs |= {'LockToken': lock_token, 'Addresses': addresses_to_update}

    if description := args.get('description'):
        kwargs |= {'Description': description}

    response = client.update_ip_set(**kwargs)

    readable_output = f'AWS Waf ip set with id {args.get("id", "")} was updated successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def list_ip_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to get a list of all IP sets"""
    kwargs = {
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Limit': arg_to_number(args.get('limit')) or 50
    }

    if next_marker := args.get('next_token'):
        kwargs |= {'NextMarker': next_marker}

    response = client.list_ip_sets(**kwargs)
    ip_sets = response.get('IPSets', [])
    readable_output = tableToMarkdown('List IP Sets',
                                      ip_sets,
                                      headers=['Name', 'Id', 'ARN', 'Description'],
                                      is_auto_json_transform=True)
    outputs = {f'{OUTPUT_PREFIX}.IpSet(val.Id === obj.Id)': ip_sets,
               f'{OUTPUT_PREFIX}(true)': {'IpSetNextToken': response.get('NextMarker', '')}}

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_key_field='Id')


def delete_ip_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to delete a specific IP set"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    get_response = client.get_ip_set(**kwargs)

    kwargs |= {'LockToken': get_response.get('LockToken', '')}

    response = client.delete_ip_set(**kwargs)

    readable_output = f'AWS Waf ip set with id {args.get("id", "")} was deleted successfully'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def create_regex_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to create a regex set"""
    tag_keys = argToList(args.get('tag_key')) or []
    tag_values = argToList(args.get('tag_value')) or []
    regex_patterns = argToList(args.get('regex_pattern')) or []
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'RegularExpressionList': build_regex_pattern_object(regex_patterns)
    }

    if description := args.get('description'):
        kwargs |= {'Description': description}
    if tags := get_tags_dict_from_args(tag_keys, tag_values):
        kwargs |= {'Tags': tags}

    response = client.create_regex_pattern_set(**kwargs)
    outputs = response.get('Summary', {})

    readable_output = f'AWS Waf regex set with id {outputs.get("Id", "")} was created successfully'

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_prefix=f'{OUTPUT_PREFIX}.RegexSet',
                          outputs_key_field='Id')


def get_regex_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to get a specific regex set"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    response = client.get_regex_pattern_set(**kwargs)

    outputs = response.get('RegexPatternSet', {})

    readable_output = tableToMarkdown('Regex Set', outputs)

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_prefix=f'{OUTPUT_PREFIX}.RegexSet',
                          outputs_key_field='Id')


def update_regex_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to update a specific regex set"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    patterns_to_update = build_regex_pattern_object(argToList(args.get('regex_pattern')))
    overwrite = argToBoolean(args.get('is_overwrite')) or False

    get_response = client.get_regex_pattern_set(**kwargs)

    lock_token = get_response.get('LockToken', '')
    original_patterns = get_response.get('RegexPatternSet', {}).get('RegularExpressionList')
    if not overwrite:
        patterns_to_update.extend(original_patterns)

    kwargs |= {'LockToken': lock_token, 'RegularExpressionList': patterns_to_update}

    if description := args.get('description'):
        kwargs |= {'Description': description}

    response = client.update_regex_pattern_set(**kwargs)

    readable_output = f'AWS Waf ip set with id {args.get("Id", "")} was updated successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def list_regex_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to get a list of all regex sets"""
    kwargs = {
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Limit': arg_to_number(args.get('limit')) or 50
    }

    if next_marker := args.get('next_token'):
        kwargs |= {'NextMarker': next_marker}

    response = client.list_regex_pattern_sets(**kwargs)
    regex_patterns = response.get('RegexPatternSets', [])
    readable_output = tableToMarkdown('List regex Sets',
                                      regex_patterns,
                                      headers=['Name', 'Id', 'ARN', 'Description'],
                                      is_auto_json_transform=True)
    outputs = {f'{OUTPUT_PREFIX}.RegexSet(val.Id === obj.Id)': regex_patterns,
               f'{OUTPUT_PREFIX}(true)': {'RegexSetNextToken': response.get('NextMarker', '')}}

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_key_field='Id')


def delete_regex_set_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to delete a specific regex set"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    get_response = client.get_regex_pattern_set(**kwargs)

    kwargs |= {'LockToken': get_response.get('LockToken', '')}

    response = client.delete_regex_pattern_set(**kwargs)

    readable_output = f'AWS Waf regex set with id {args.get("id", "")} was deleted successfully'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def list_rule_group_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to get a list of all rule groups"""
    kwargs = {
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Limit': arg_to_number(args.get('limit')) or 50
    }

    if next_marker := args.get('next_token'):
        kwargs |= {'NextMarker': next_marker}

    response = client.list_rule_groups(**kwargs)
    rule_groups = response.get('RuleGroups', [])
    outputs = {f'{OUTPUT_PREFIX}.RuleGroup(val.Id === obj.Id)': rule_groups,
               f'{OUTPUT_PREFIX}(true)': {'RuleGroupNextToken': response.get('NextMarker', '')}}
    readable_output = tableToMarkdown('List rule groups',
                                      rule_groups,
                                      headers=['Name', 'Id', 'ARN', 'Description'],
                                      is_auto_json_transform=True)

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_key_field='Id')


def get_rule_group_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to get a specific rule group"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    response = client.get_rule_group(**kwargs)
    response = convert_dict_values_bytes_to_str(response)
    outputs = response.get('RuleGroup', {})
    readable_output = tableToMarkdown('Rule group', outputs, headers=['Id', 'Name', 'Description'])

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_prefix=f'{OUTPUT_PREFIX}.RuleGroup',
                          outputs_key_field='Id')


def delete_rule_group_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to delete a specific rule group"""
    kwargs = {
        'Name': args.get('name', ''),
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Id': args.get('id', '')
    }

    get_response = client.get_rule_group(**kwargs)

    kwargs |= {'LockToken': get_response.get('LockToken', '')}

    response = client.delete_rule_group(**kwargs)

    readable_output = f'AWS Waf rule group with id {args.get("id", "")} was deleted successfully'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def create_rule_group_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to create a rule group"""
    tag_keys = argToList(args.get('tag_key')) or []
    tag_values = argToList(args.get('tag_value')) or []
    name = args.get('name', '')
    cloud_watch_metrics_enabled = argToBoolean(args.get('cloud_watch_metrics_enabled', '')) or True
    metric_name = args.get('metric_name', '') or name
    sampled_requests_enabled = argToBoolean(args.get('sampled_requests_enabled', '')) or True

    kwargs = {
        'Name': name,
        'Scope': SCOPE_MAP[args.get('scope') or DEFAULT_SCOPE],
        'Capacity': arg_to_number(args.get('capacity', '')),
        'VisibilityConfig': build_visibility_config_object(cloud_watch_metrics_enabled=cloud_watch_metrics_enabled,
                                                           metric_name=metric_name,
                                                           sampled_requests_enabled=sampled_requests_enabled)
    }

    if description := args.get('description'):
        kwargs |= {'Description': description}
    if tags := get_tags_dict_from_args(tag_keys, tag_values):
        kwargs |= {'Tags': tags}

    response = client.create_rule_group(**kwargs)
    outputs = response.get('Summary', {})

    readable_output = f'AWS Waf rule group with id {outputs.get("Id", "")} was created successfully'

    return CommandResults(readable_output=readable_output,
                          outputs=outputs,
                          raw_response=response,
                          outputs_prefix=f'{OUTPUT_PREFIX}.RuleGroup',
                          outputs_key_field='Id')


def create_ip_rule_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to create an ip rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)

    rule = build_new_rule_object(args, rule_group_visibility_config, build_ip_rule_object)
    updated_rules = append_new_rule(rules, rule)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf ip rule with name {args.get("rule_name", "")} was created successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def create_country_rule_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to create a country rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)

    rule = build_new_rule_object(args, rule_group_visibility_config, build_country_rule_object)
    updated_rules = append_new_rule(rules, rule)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf country rule with name {args.get("rule_name", "")} was created successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def create_string_match_rule_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to create a string match rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)

    rule = build_new_rule_object(args, rule_group_visibility_config, build_string_match_rule_object)
    updated_rules = append_new_rule(rules, rule)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf string match rule with name {args.get("rule_name", "")} was created successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def delete_rule_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to delete a specific rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)
    rule_name = args.get('rule_name', '')
    updated_rules = delete_rule(rule_name, rules)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf rule with id {args.get("Id", "")} was deleted successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def add_ip_statement_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to add an ip statement to a rule"""
    kwargs = get_required_args_for_get_rule_group(args)
    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)
    ip_set_arns = argToList(args.get('ip_set_arn'))
    statements = [build_ip_statement(ip_set_arn) for ip_set_arn in ip_set_arns]
    updated_rules = create_rules_list_with_new_rule_statement(args, statements, rules)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf ip statement was added to rule with name {args.get("rule_name", "")} successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def add_country_statement_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to add a country statement to a rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)
    country_codes = argToList(args.get('country_codes')) or []
    statement = [build_country_statement(country_codes)]
    updated_rules = create_rules_list_with_new_rule_statement(args, statement, rules)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf country statement was added to rule with name {args.get("rule_name", "")} ' \
                      f'successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def add_string_match_statement_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to add a string match statement to a rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)

    statement = [build_string_match_statement(**args)]
    updated_rules = create_rules_list_with_new_rule_statement(args, statement, rules)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf string match statement was added to rule with name {args.get("rule_name", "")} ' \
                      f'successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def add_json_statement_command(client: boto3.client, args: dict) -> CommandResults:
    """ Command to add a json object represents a statement to a rule"""
    kwargs = get_required_args_for_get_rule_group(args)

    rules, rule_group_visibility_config, lock_token = get_required_response_fields_from_rule_group(client, kwargs)

    statement = json.loads(args.get('statement_json') or '{}')
    updated_rules = create_rules_list_with_new_rule_statement(args, statement, rules)

    response = update_rule_group_rules(client=client,
                                       kwargs=kwargs,
                                       lock_token=lock_token,
                                       updated_rules=updated_rules,
                                       rule_group_visibility_config=rule_group_visibility_config)

    readable_output = f'AWS Waf json statement was added to rule with name {args.get("rule_name", "")} successfully.'

    return CommandResults(readable_output=readable_output,
                          raw_response=response)


def template_json_command(args: dict) -> CommandResults:  # pragma: no cover
    """ Command to get a json template represents a statement"""
    statement_type = args.get('statement_type', '')
    web_request_component = args.get('web_request_component', '')
    if statement_type == 'Ip Set':
        readable_output = json.dumps(build_ip_statement(ip_set_arn='The Ip Set ARN'))

    elif statement_type == 'Country':
        readable_output = json.dumps(build_country_statement(country_codes=['country code1, country code2...']))

    else:
        if not web_request_component:
            raise DemistoException('Please provide web_request_component for string match and regex match ')
        web_request_component = WEB_REQUEST_COMPONENT_MAP[web_request_component]
        if statement_type == 'String Match':
            output = build_byte_match_statement(web_request_component=web_request_component,
                                                oversize_handling='CONTINUE | MATCH | NO_MATCH',
                                                text_transformation=TEXT_TRANSFORMATIONS,
                                                string_to_match='The string to match',
                                                match_type='all')
            match_statement = BYTE_MATCH_STATEMENT

        else:  # statement_type == 'Regex Pattern':
            output = build_regex_match_statement(web_request_component=web_request_component,
                                                 oversize_handling='CONTINUE | MATCH | NO_MATCH',
                                                 text_transformation=TEXT_TRANSFORMATIONS,
                                                 regex_set_arn="The regex set ARN")
            match_statement = REGEX_MATCH_STATEMENT

        statement = {match_statement: output}
        readable_output = json.dumps(statement)

    return CommandResults(readable_output=readable_output)


''' MAIN FUNCTION '''


def main() -> None:  # pragma: no cover
    params = demisto.params()
    aws_default_region = params.get('defaultRegion')
    aws_role_arn = params.get('roleArn')
    aws_role_session_name = params.get('roleSessionName')
    aws_role_session_duration = params.get('sessionDuration')
    aws_access_key_id = params.get('access_key', {}).get('password') or params.get('access_key')
    aws_secret_access_key = params.get('secret_key', {}).get('password') or params.get('secret_key')
    verify_certificate = not params.get('insecure', True)
    timeout = params.get('timeout') or 1
    retries = params.get('retries') or 5

    try:
        validate_params(aws_default_region, aws_role_arn, aws_role_session_name, aws_access_key_id,
                        aws_secret_access_key)

        aws_client = AWSClient(aws_default_region=aws_default_region,
                               aws_role_arn=aws_role_arn,
                               aws_role_session_name=aws_role_session_name,
                               aws_role_policy=None,
                               aws_role_session_duration=aws_role_session_duration,
                               aws_access_key_id=aws_access_key_id,
                               aws_secret_access_key=aws_secret_access_key,
                               verify_certificate=verify_certificate,
                               timeout=timeout,
                               retries=retries)
        args = demisto.args()
        command = demisto.command()
        client = aws_client.aws_session(service=SERVICE, region=args.get('region'))

        if command == 'test-module':
            # This is the call made when pressing the integration test button.
            return_results(connection_test(client))

        elif command == 'aws-waf-ip-set-create':
            result = create_ip_set_command(client, args)
        elif command == 'aws-waf-ip-set-get':
            result = get_ip_set_command(client, args)
        elif command == 'aws-waf-ip-set-update':
            result = update_ip_set_command(client, args)
        elif command == 'aws-waf-ip-set-list':
            result = list_ip_set_command(client, args)
        elif command == 'aws-waf-ip-set-delete':
            result = delete_ip_set_command(client, args)

        elif command == 'aws-waf-regex-set-create':
            result = create_regex_set_command(client, args)
        elif command == 'aws-waf-regex-set-get':
            result = get_regex_set_command(client, args)
        elif command == 'aws-waf-regex-set-update':
            result = update_regex_set_command(client, args)
        elif command == 'aws-waf-regex-set-list':
            result = list_regex_set_command(client, args)
        elif command == 'aws-waf-regex-set-delete':
            result = delete_regex_set_command(client, args)

        elif command == 'aws-waf-rule-group-list':
            result = list_rule_group_command(client, args)
        elif command == 'aws-waf-rule-group-get':
            result = get_rule_group_command(client, args)
        elif command == 'aws-waf-rule-group-delete':
            result = delete_rule_group_command(client, args)
        elif command == 'aws-waf-rule-group-create':
            result = create_rule_group_command(client, args)

        elif command == 'aws-waf-ip-rule-create':
            result = create_ip_rule_command(client, args)
        elif command == 'aws-waf-country-rule-create':
            result = create_country_rule_command(client, args)
        elif command == 'aws-waf-string-match-rule-create':
            result = create_string_match_rule_command(client, args)
        elif command == 'aws-waf-rule-delete':
            result = delete_rule_command(client, args)

        elif command == 'aws-waf-ip-statement-add':
            result = add_ip_statement_command(client, args)
        elif command == 'aws-waf-country-statement-add':
            result = add_country_statement_command(client, args)
        elif command == 'aws-waf-string-match-statement-add':
            result = add_string_match_statement_command(client, args)
        elif command == 'aws-waf-statement-json-add':
            result = add_json_statement_command(client, args)

        elif command == 'aws-waf-statement-json-template-get':
            result = template_json_command(args)

        else:
            raise NotImplementedError(f'Command {command} is not implemented in AWS WAF integration.')

        return_results(result)

    # Log exceptions and return errors
    except Exception as e:
        return_error(f'Failed to execute {command} command.\nError:\n{str(e)}')


''' ENTRY POINT '''

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