/*
 *
 * (C) 2015-23 - ntop.org
 *
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 */

#include "ntop_includes.h"

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

TimelineExtract::TimelineExtract() {
  extraction.id = 0;
  status_code = 0;
  running = false;
  shutdown = false;
}

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

TimelineExtract::~TimelineExtract() { stop(); }

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

#ifdef HAVE_PF_RING
pfring *TimelineExtract::openTimeline(const char *timeline_path, time_t from,
                                      time_t to, const char *bpf_filter) {
  char from_buff[24], to_buff[24];
  pfring *handle = NULL;
  char *filter;
  struct tm *time_info;
  int rc, len;
  char timeline_ifname[MAX_PATH];

  snprintf(timeline_ifname, sizeof(timeline_ifname), "timeline:%s", timeline_path);

  handle = pfring_open(timeline_ifname, 16384, 0);

  if (handle == NULL) {
    ntop->getTrace()->traceEvent(TRACE_ERROR, "Unable to open %s",
                                 timeline_path);
    status_code = 4; /* Unable to open timeline */
    goto error;
  }

  len = 64 + (bpf_filter ? strlen(bpf_filter) : 0);
  filter = (char *)malloc(len);

  if (filter == NULL) {
    ntop->getTrace()->traceEvent(TRACE_ERROR, "Unable to allocate memory");
    status_code = 3; /* Memory allocation failure */
    goto close_pfring;
  }

  filter[0] = '\0';

  time_info = localtime(&from);
  strftime(from_buff, sizeof(from_buff), "%Y-%m-%d %H:%M:%S", time_info);

  time_info = localtime(&to);
  strftime(to_buff, sizeof(to_buff), "%Y-%m-%d %H:%M:%S", time_info);

  snprintf(filter, len, "start %s and end %s", from_buff, to_buff);

  if (bpf_filter && strlen(bpf_filter) > 0)
    sprintf(&filter[strlen(filter)], " and %s", bpf_filter);

  ntop->getTrace()->traceEvent(
      TRACE_INFO, "Running extraction from '%s' matching filter '%s'",
      timeline_path, filter);

  rc = pfring_set_bpf_filter(handle, filter);

  free(filter);

  if (rc != 0) {
    ntop->getTrace()->traceEvent(TRACE_ERROR, "Unable to set filter '%s' (%d)",
                                 filter, rc);
    status_code = 5; /* Unable to set filter */
    goto close_pfring;
  }

  if (pfring_enable_ring(handle) != 0) {
    ntop->getTrace()->traceEvent(
        TRACE_ERROR, "Unable to start extraction on %s", timeline_path);
    status_code = 6; /* Unable to open timeline */
    goto close_pfring;
  }

  return handle;

close_pfring:
  pfring_close(handle);

error:
  return NULL;
}

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

pfring *TimelineExtract::openTimelineFromInterface(NetworkInterface *iface,
                                                   time_t from, time_t to,
                                                   const char *bpf_filter) {
  char timeline_path[MAX_PATH];

  snprintf(timeline_path, sizeof(timeline_path), "%s/%d/timeline",
           ntop->getPrefs()->get_pcap_dir(), iface->get_id());

  return openTimeline(timeline_path, from, to, bpf_filter);
}

#endif

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

bool TimelineExtract::extractToDisk(u_int32_t id, NetworkInterface *iface,
                                    time_t from, time_t to,
                                    const char *bpf_filter, u_int64_t max_bytes,
                                    const char *timeline_path) {
  bool completed = false;
#ifdef HAVE_PF_RING
  char out_path[MAX_PATH];
  PacketDumper *dumper;
  pfring *handle;
  u_char *packet = NULL;
  struct pfring_pkthdr header;
  struct pcap_pkthdr *h;

  memset(&header, 0, sizeof(header));

  shutdown = false;
  stats.packets = stats.bytes = 0;
  status_code = 1; /* default: unexpected error */

  snprintf(out_path, sizeof(out_path), "%s/%u/extr_pcap/%u",
           ntop->getPrefs()->get_pcap_dir(), iface->get_id(), id);

  dumper = new (std::nothrow) PacketDumper(iface, out_path);

  if (dumper == NULL) {
    ntop->getTrace()->traceEvent(TRACE_ERROR,
                                 "Unable to initialize packet dumper");
    status_code = 2; /* Unable to initialize dumper */
    goto error;
  }

  if (!timeline_path || timeline_path[0] == '\0')
    handle = openTimelineFromInterface(iface, from, to, bpf_filter);
  else
    handle = openTimeline(timeline_path, from, to, bpf_filter);

  if (handle == NULL) {
    goto delete_dumper;
  }

  ntop->getTrace()->traceEvent(TRACE_INFO, "Dumping traffic to '%s'", out_path);

  while (!shutdown && !ntop->getGlobals()->isShutdown() &&
         pfring_recv(handle, &packet, 0, &header, 0) > 0) {
    h = (struct pcap_pkthdr *)&header;
    dumper->dumpPacket(h, packet);
    stats.packets++;
    stats.bytes += sizeof(struct pcap_disk_pkthdr) + h->caplen;
    if (max_bytes != 0 && stats.bytes >= max_bytes) break;
  }

  status_code = 0; /* Successfully completed */
  completed = true;

  pfring_close(handle);

delete_dumper:
  delete dumper;

error:
#else
  status_code = 7; /* No PF_RING support */
#endif

  ntop->getTrace()->traceEvent(TRACE_INFO, "Extraction #%u %s", id,
                               completed ? "completed" : "failed");

  return completed;
}

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

bool TimelineExtract::extractLive(struct mg_connection *conn,
                                  NetworkInterface *iface, time_t from,
                                  time_t to, const char *bpf_filter,
                                  const char *timeline_path) {
  bool completed = false;
#ifdef HAVE_PF_RING
  pfring *handle;
  u_char *packet = NULL;
  struct pfring_pkthdr h;
  struct pcap_file_header pcaphdr;
  struct pcap_disk_pkthdr pkthdr;
  bool http_client_disconnected = false;
  int rc;

  memset(&h, 0, sizeof(h));
  stats.packets = stats.bytes = 0;

  ntop->getTrace()->traceEvent(TRACE_INFO, "Running live extraction");

  if (!timeline_path || timeline_path[0] == '\0')
    handle = openTimelineFromInterface(iface, from, to, bpf_filter);
  else
    handle = openTimeline(timeline_path, from, to, bpf_filter);

  if (handle == NULL) {
    goto error;
  }

  Utils::init_pcap_header(&pcaphdr, pfring_get_link_type(handle),
                          pfring_get_caplen(handle));

  if (!Utils::mg_write_retry(conn, (u_char *)&pcaphdr, sizeof(pcaphdr)))
    http_client_disconnected = true;

  while (!http_client_disconnected && !ntop->getGlobals()->isShutdown() &&
         (rc = pfring_recv(handle, &packet, 0, &h, 0)) > 0) {
    pkthdr.ts.tv_sec = h.ts.tv_sec;
    pkthdr.ts.tv_usec = h.ts.tv_usec, pkthdr.caplen = h.caplen;
    pkthdr.len = h.len;

    if (!Utils::mg_write_retry(conn, (u_char *)&pkthdr, sizeof(pkthdr)) ||
        !Utils::mg_write_retry(conn, (u_char *)packet, h.caplen))
      http_client_disconnected = true;

    stats.packets++;
    stats.bytes += sizeof(struct pcap_disk_pkthdr) + h.caplen;
  }

  completed = true;
  pfring_close(handle);

error:
  if (completed)
    ntop->getTrace()->traceEvent(
        TRACE_INFO, "Live extraction %s %s", "completed",
        http_client_disconnected ? "(disconnected)" : "");
  else
    ntop->getTrace()->traceEvent(
        TRACE_ERROR, "Live extraction %s %s", "failed",
        http_client_disconnected ? "(disconnected)" : "");
#endif

  return completed;
}

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

static void *extractionThread(void *ptr) {
  TimelineExtract *extr = (TimelineExtract *)ptr;
  Utils::setThreadName("ntopng-extract");

  extr->extractToDisk(extr->getID(), extr->getNetworkInterface(),
                      extr->getFrom(), extr->getTo(), extr->getFilter(),
                      extr->getMaxBytes(), extr->getTimelinePath());

  extr->cleanupJob();
  return NULL;
}

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

void TimelineExtract::runExtractionJob(u_int32_t id, NetworkInterface *iface,
                                       time_t from, time_t to,
                                       const char *bpf_filter,
                                       u_int64_t max_bytes,
                                       const char *timeline_path) {
  running = true;

  extraction.id = id;
  extraction.iface = iface;
  extraction.from = from;
  extraction.to = to;
  extraction.bpf_filter = strdup(bpf_filter);
  extraction.max_bytes = max_bytes;
  extraction.timeline_path = timeline_path;

  pthread_create(&extraction_thread, NULL, extractionThread, (void *)this);
}

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

void TimelineExtract::stopExtractionJob(u_int32_t id) {
  if (running && extraction.id == id) stop();
}

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

void TimelineExtract::stop() {
  void *res;

  shutdown = true;

  if (running) pthread_join(extraction_thread, &res);
}

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

void TimelineExtract::cleanupJob() {
  if (extraction.bpf_filter) free(extraction.bpf_filter);

  running = false;
}

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

void TimelineExtract::getStatus(lua_State *vm) {
  lua_newtable(vm);

  if (extraction.id) {
    lua_newtable(vm);

    lua_push_uint64_table_entry(vm, "id", extraction.id);
    lua_push_uint64_table_entry(vm, "extracted_pkts", stats.packets);
    lua_push_uint64_table_entry(vm, "extracted_bytes", stats.bytes);
    lua_push_uint64_table_entry(vm, "status", status_code);

    lua_pushinteger(vm, extraction.id);
    lua_insert(vm, -2);
    lua_settable(vm, -3);
  }
}

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