/*
 * ndpi_geoip.c
 *
 * Copyright (C) 2021 - ntop.org
 *
 * This file is part of nDPI, an open source deep packet inspection
 * library.
 *
 * nDPI is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * nDPI is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with nDPI.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <stdlib.h>
#include <errno.h>
#include <sys/types.h>
#include <stdint.h>

#include "ndpi_api.h"
#include "ndpi_private.h"
#include "ndpi_config.h"

#ifdef HAVE_MAXMINDDB
#include <maxminddb.h>
#endif

/* ********************************************************************************* */

int ndpi_load_geoip(struct ndpi_detection_module_struct *ndpi_str,
                    const char *ip_city_data, const char *ip_as_data)
{
#ifdef HAVE_MAXMINDDB
  int status;

  ndpi_str->mmdb_city = (void *)ndpi_malloc(sizeof(MMDB_s));
  ndpi_str->mmdb_as = (void *)ndpi_malloc(sizeof(MMDB_s));

  if ((ndpi_str->mmdb_city == NULL) || (ndpi_str->mmdb_as == NULL))
    return (-1);

  /* Open the MMDB files */
  if ((status = MMDB_open(ip_city_data, MMDB_MODE_MMAP, (MMDB_s *)ndpi_str->mmdb_city)) != MMDB_SUCCESS)
    return (-1);
  else
    ndpi_str->mmdb_city_loaded = 1;

  if ((status = MMDB_open(ip_as_data, MMDB_MODE_MMAP, (MMDB_s *)ndpi_str->mmdb_as)) != MMDB_SUCCESS)
    return (-2);
  else
    ndpi_str->mmdb_as_loaded = 1;

  return (0);
#else
  (void)ndpi_str;
  (void)ip_city_data;
  (void)ip_as_data;
  return (-3);
#endif
}

/* ********************************************************************************* */

void ndpi_free_geoip(struct ndpi_detection_module_struct *ndpi_str)
{
#ifdef HAVE_MAXMINDDB
  if (ndpi_str->mmdb_city_loaded)
    MMDB_close((MMDB_s *)ndpi_str->mmdb_city);
  if (ndpi_str->mmdb_as_loaded)
    MMDB_close((MMDB_s *)ndpi_str->mmdb_as);

  ndpi_free(ndpi_str->mmdb_city);
  ndpi_free(ndpi_str->mmdb_as);
#else
  (void)ndpi_str;
#endif
}

/* ********************************************************************************* */

int ndpi_get_geoip_asn(struct ndpi_detection_module_struct *ndpi_str, char *ip, u_int32_t *asn)
{
#ifdef HAVE_MAXMINDDB
  int gai_error, mmdb_error, status;
  MMDB_lookup_result_s result;
  MMDB_entry_data_s entry_data;

  if (ndpi_str->mmdb_as_loaded)
  {
    result = MMDB_lookup_string((MMDB_s *)ndpi_str->mmdb_as, ip, &gai_error, &mmdb_error);

    if ((gai_error != 0) || (mmdb_error != MMDB_SUCCESS) || (!result.found_entry))
      *asn = 0;
    else
    {
      /* Get the ASN */
      if ((status = MMDB_get_value(&result.entry, &entry_data, "autonomous_system_number", NULL)) == MMDB_SUCCESS)
      {
        if (entry_data.has_data && entry_data.type == MMDB_DATA_TYPE_UINT32)
          *asn = entry_data.uint32;
        else
          *asn = 0;
      }
    }

    return (0);
  }
#else
  (void)ndpi_str;
  (void)ip;
  (void)asn;
#endif

  return (-2);
}

int ndpi_get_geoip_aso(struct ndpi_detection_module_struct *ndpi_str, char *ip, char *aso, u_int8_t aso_len)
{
#ifdef HAVE_MAXMINDDB
  int gai_error, mmdb_error, status;
  MMDB_lookup_result_s result;
  MMDB_entry_data_s entry_data;

  if (ndpi_str->mmdb_as_loaded && aso_len > 0)
  {
    result = MMDB_lookup_string((MMDB_s *)ndpi_str->mmdb_as, ip, &gai_error, &mmdb_error);

    if ((gai_error != 0) || (mmdb_error != MMDB_SUCCESS) || (!result.found_entry))
      aso[0] = '\0';
    else
    {
      /* Get the ASO */
      if (aso_len > 0)
      {
        status = MMDB_get_value(&result.entry, &entry_data, "autonomous_system_organization", NULL);
        if (status != MMDB_SUCCESS || !entry_data.has_data)
          aso[0] = '\0';
        else
        {
          int str_len = ndpi_min(entry_data.data_size, aso_len);

          memcpy(aso, entry_data.utf8_string, str_len);
          aso[str_len] = '\0';
        }
      }
    }

    return (0);
  }
#else
  (void)ndpi_str;
  (void)ip;
  (void)aso;
  (void)aso_len;
#endif

  return (-2);
}

/* ********************************************************************************* */

int ndpi_get_geoip_country_continent(struct ndpi_detection_module_struct *ndpi_str, char *ip,
                                     char *country_code, u_int8_t country_code_len,
                                     char *continent, u_int8_t continent_len)
{
#ifdef HAVE_MAXMINDDB
  int gai_error, mmdb_error;
  MMDB_lookup_result_s result;
  MMDB_entry_data_s entry_data;

  if (ndpi_str->mmdb_city_loaded)
  {
    int status;

    result = MMDB_lookup_string((MMDB_s *)ndpi_str->mmdb_city, ip, &gai_error, &mmdb_error);

    if ((gai_error != 0) || (mmdb_error != MMDB_SUCCESS) || (!result.found_entry))
      country_code[0] = '\0';
    else
    {
      if (country_code_len > 0)
      {
        status = MMDB_get_value(&result.entry, &entry_data, "country", "iso_code", NULL);

        if ((status != MMDB_SUCCESS) || (!entry_data.has_data))
          country_code[0] = '\0';
        else
        {
          int str_len = ndpi_min(entry_data.data_size, country_code_len);

          memcpy(country_code, entry_data.utf8_string, str_len);
          country_code[str_len] = '\0';
        }
      }

      if (continent_len > 0)
      {
        status = MMDB_get_value(&result.entry, &entry_data, "continent", "names", "en", NULL);

        if ((status != MMDB_SUCCESS) || (!entry_data.has_data))
          continent[0] = '\0';
        else
        {
          int str_len = ndpi_min(entry_data.data_size, continent_len);

          memcpy(continent, entry_data.utf8_string, str_len);
          continent[str_len] = '\0';
        }
      }
    }

    return (0);
  }
#else
  (void)ndpi_str;
  (void)ip;
  (void)country_code;
  (void)country_code_len;
  (void)continent;
  (void)continent_len;
#endif

  return (-2);
}

int ndpi_get_geoip_country_continent_city(struct ndpi_detection_module_struct *ndpi_str, char *ip,
                                          char *country_code, u_int8_t country_code_len,
                                          char *continent, u_int8_t continent_len,
                                          char *city, u_int8_t city_len)
{
#ifdef HAVE_MAXMINDDB
  int gai_error, mmdb_error;
  MMDB_lookup_result_s result;
  MMDB_entry_data_s entry_data;

  if (ndpi_str->mmdb_city_loaded)
  {
    int status;

    result = MMDB_lookup_string((MMDB_s *)ndpi_str->mmdb_city, ip, &gai_error, &mmdb_error);

    if ((gai_error != 0) || (mmdb_error != MMDB_SUCCESS) || (!result.found_entry))
      country_code[0] = '\0';
    else
    {
      if (country_code_len > 0)
      {
        status = MMDB_get_value(&result.entry, &entry_data, "country", "iso_code", NULL);

        if ((status != MMDB_SUCCESS) || (!entry_data.has_data))
          country_code[0] = '\0';
        else
        {
          int str_len = ndpi_min(entry_data.data_size, country_code_len);

          memcpy(country_code, entry_data.utf8_string, str_len);
          country_code[str_len] = '\0';
        }
      }

      if (continent_len > 0)
      {
        status = MMDB_get_value(&result.entry, &entry_data, "continent", "names", "en", NULL);

        if ((status != MMDB_SUCCESS) || (!entry_data.has_data))
          continent[0] = '\0';
        else
        {
          int str_len = ndpi_min(entry_data.data_size, continent_len);

          memcpy(continent, entry_data.utf8_string, str_len);
          continent[str_len] = '\0';
        }
      }

      if (city_len > 0)
      {
        status = MMDB_get_value(&result.entry, &entry_data, "city", "names", "en", NULL);

        if ((status != MMDB_SUCCESS) || (!entry_data.has_data))
          city[0] = '\0';
        else
        {
          int len = ndpi_min(entry_data.data_size, city_len);

          memcpy(city, entry_data.utf8_string, len);
          city[len] = 0;
        }
      }

      return (0);
    }
  }
#else
  (void)ndpi_str;
  (void)ip;
  (void)country_code;
  (void)country_code_len;
  (void)continent;
  (void)continent_len;
  (void)city;
  (void)city_len;
#endif

  return (-2);
}