from time import sleep

from CommonServerPython import *
from JSONFeedApiModule import *

INTEGRATION_NAME = "MalwareBazaar Feed"


def custom_mapping_function(mapping: Dict, indicator: Dict, attributes: Dict):
    for map_key in mapping:
        if map_key in attributes:
            fields = mapping[map_key].split(".")
            if len(fields) > 1:
                if indicator["fields"].get(fields[0]):
                    indicator["fields"][fields[0]][0].update({fields[1]: attributes.get(map_key)})
                else:
                    indicator["fields"][fields[0]] = [{fields[1]: attributes.get(map_key)}]
            elif mapping[map_key] == "tags":
                # Merge API response tags with the existing feedTags to avoid overwriting
                # the configured feedTags value set by the JSONFeedApiModule.
                existing_tags: list = indicator["fields"].get("tags") or []
                api_tags = attributes.get(map_key) or []
                if isinstance(api_tags, list):
                    indicator["fields"]["tags"] = list(set(existing_tags + api_tags))  # Use set to remove duplicates
                else:
                    demisto.debug(f"{api_tags=} is not a list (unexpected), keeping existing feedTags unchanged.")
            else:
                indicator["fields"][mapping[map_key]] = attributes.get(map_key)
        if map_key == "url_download":
            indicator["fields"]["downloadurl"] = f"{mapping[map_key]}{indicator['value']}/"


def custom_build_relationships(feed_config: Dict, _mapping: Dict, indicator_data: Dict) -> List[dict]:
    if indicator_data.get(feed_config.get("relation_entity_b")):
        relationships_lst = EntityRelationship(
            name=feed_config.get("relation_name"),
            entity_a=indicator_data.get("value"),
            entity_a_type=indicator_data.get("type"),
            entity_b=indicator_data.get(feed_config.get("relation_entity_b")),
            entity_b_type=feed_config.get("relation_entity_b_type"),
            reverse_name=feed_config.get("reverse_relationship_name"),
        )
        return [relationships_lst.to_indicator()]
    return []


def feed_main_with_error_handling(params: dict, command: str, attempts: int):
    """A recursive function which calls the feed_main function and performs 3 retries when a known MalwareBazaar issue occurs.

    Args:
        params: the parameters of the feed.
        command: the command being called.
        attempts: the number of attempts to perform in case of a known server error.

    Raises:
        e: In case the error is not a known error, or the last retry has failed.
    """
    try:
        if not params["headers"]["Auth-Key"]:
            raise ValueError("Missing required parameter Auth Key. Please set this parameter in the instance configuration.")

        feed_main(params, "MalwareBazaar Feed", "malwarebazzar")
    except Exception as e:
        # Handles known server side and network exceptions
        msg = str(e)
        if (
            command == "fetch-indicators"
            and attempts > 0
            and (
                ("502 Server Error: Bad Gateway") in msg
                or "Connection broken: IncompleteRead" in msg
                or "503 Server Error: Service Unavailable" in msg
            )
        ):
            demisto.error(f"An Error Occurred during the run of the 'fetch-indicators' command: {msg}.")
            demisto.error(f"This is attempt number: {4 - attempts}.\nRetrying in 10 seconds...")
            sleep(10)
            attempts -= 1
            feed_main_with_error_handling(params, command, attempts)
        else:
            raise e


def main() -> None:  # pragma: no cover
    api_key = demisto.params().get("credentials", {}).get("password")
    params = {k: v for k, v in demisto.params().items() if v is not None}

    params["headers"] = {"Auth-Key": api_key}

    params["indicator_type"] = FeedIndicatorType.File
    params["feed_name_to_config"] = {
        "File": {
            "url": f'{params.get("url")}/api/v1/',
            "extractor": "data",
            "indicator": "sha256_hash",
            "indicator_type": FeedIndicatorType.File,
            "relation_name": EntityRelationship.Relationships.INDICATOR_OF,
            "reverse_relationship_name": EntityRelationship.Relationships.INDICATED_BY,
            "relation_entity_b": "signature",
            "relation_entity_b_type": "Malware",
            "create_relations_function": custom_build_relationships,
            "mapping_function": custom_mapping_function,
            "mapping": {
                "sha256_hash": "sha256",
                "sha1_hash": "sha1",
                "md5_hash": "md5",
                "first_seen": "first_seen_by_source",
                "last_seen": "last_seen_by_source",
                "file_name": "Associated File Names",
                "file_size": "size",
                "file_type": "filetype",
                "reporter": "reported_by",
                "imphash": "imphash",
                "ssdeep": "ssdeep",
                "tags": "tags",
                "url_download": "https://bazaar.abuse.ch/sample/",
            },
        }
    }
    params["data"] = {
        "query": "get_recent",
        "selector": "time",
    }  # query params to get only the recent changes for the incremental feed

    command = demisto.command()

    attempts = 3  # Due to multiple server side errors from malwareBazaar we perform 3 retries when needed

    feed_main_with_error_handling(params, command, attempts)


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