IP2LocationIO

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


import urllib3
from typing import Any

urllib3.disable_warnings()


class Client(BaseClient):
    def get_ip_geolocation(self, ip: str, api_key: str) -> dict[str, Any]:
        return self._http_request(method="GET", url_suffix="/", params={"ip": ip, "key": api_key})


def test_module(client: Client) -> str:
    try:
        client._http_request(method="GET", url_suffix="/", params={"ip": "8.8.8.8"})
    except DemistoException as e:
        if "Forbidden" in str(e):
            return "Authorization Error: make sure API Key is correctly set"
        else:
            raise e

    return "ok"


def ip_geolocation_command(
    client: Client, args: dict[str, Any], reliability: DBotScoreReliability, api_key: str
) -> list[CommandResults]:
    ips = argToList(args.get("ip"))
    if len(ips) == 0:
        raise ValueError("IP(s) not specified")

    command_results: list[CommandResults] = []

    for ip in ips:
        if not is_ip_valid(ip, accept_v6_ips=True):  # check IP's validity
            raise ValueError(f'IP "{ip}" is not valid')
        ip_data = client.get_ip_geolocation(ip, api_key)
        ip_data["ip"] = ip

        dbot_score = Common.DBotScore(
            indicator=ip,
            indicator_type=DBotScoreType.IP,
            integration_name="IP2LocationIO",
            score=Common.DBotScore.NONE,
            reliability=reliability,
        )

        ip_standard_context = Common.IP(
            ip=ip,
            geo_country=ip_data.get("country_name"),
            geo_latitude=ip_data.get("latitude"),
            geo_longitude=ip_data.get("longitude"),
            geo_description=f"{ip_data.get('city_name')}, {ip_data.get('region_name')}, {ip_data.get('country_name')}",
            region=ip_data.get("region"),
            asn=f"AS{ip_data.get('asn')}",
            dbot_score=dbot_score,
        )

        ip_context_excluded_fields = ["objects", "nir"]
        ip_data = {k: ip_data[k] for k in ip_data if k not in ip_context_excluded_fields}

        readable_output = tableToMarkdown("IP", ip_data)

        command_results.append(
            CommandResults(
                readable_output=readable_output,
                outputs_prefix="IP2LocationIO.IP",
                outputs_key_field="ip",
                outputs=ip_data,
                indicator=ip_standard_context,
            )
        )
    return command_results


def main() -> None:
    params = demisto.params()
    args = demisto.args()
    command = demisto.command()

    api_key = params.get("credentials", {}).get("password")

    base_url = urljoin(params.get("url").rstrip("/"), "")

    verify_certificate = not params.get("insecure", False)

    proxy = params.get("proxy", False)

    reliability = params.get("integrationReliability", DBotScoreReliability.C)

    demisto.debug(f"Command being called is {command}")
    try:
        client = Client(base_url=base_url, verify=verify_certificate, proxy=proxy)

        if command == "test-module":
            result = test_module(client)
            return_results(result)

        elif command == "ip":
            return_results(ip_geolocation_command(client, args, reliability, api_key))

        else:
            raise NotImplementedError(f"Command {command} is not implemented")

    except Exception as e:
        return_error(f"Failed to execute {command} command.\nError:\n{str(e)}")


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