/*
 * ZMap Copyright 2013 Regents of the University of Michigan
 *
 * XMap Copyright 2021 Xiang Li from Network and Information Security Lab
 * Tsinghua University
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 */

#include "shard.h"

#include <assert.h>
#include <string.h>

#include "state.h"

#include "../lib/blocklist.h"
#include "../lib/gmp-ext.h"
#include "../lib/logger.h"

static inline void shard_get_next_elem(shard_t *shard) {
    // no overflow
    mpz_mul(shard->current, shard->current, shard->params.factor);
    mpz_mod(shard->current, shard->current, shard->params.modulus);
}

static void _shard_get_current_ipvx_port_index(mpz_t ipvx, shard_t *shard) {
    mpz_t index;
    mpz_init(index);
    mpz_sub_ui(index, shard->current, 1);
    blocklist_lookup_index_for_ipvx_port_index(ipvx, index);
    mpz_clear(index);
}

static void _shard_get_next_ipvx_port_index(mpz_t ipvx, shard_t *shard) {
    if (mpz_eq_ui(shard->current, PMAP_SHARD_DONE)) {
        mpz_set_ui(ipvx, PMAP_SHARD_DONE);
        return;
    }

    mpz_t index;
    mpz_init(index);
    while (1) {
        shard_get_next_elem(shard);
        if (mpz_eq(shard->current, shard->params.last)) {
            mpz_set_ui(shard->current, PMAP_SHARD_DONE);
            shard->iterations++;
            mpz_set_ui(ipvx, PMAP_SHARD_DONE);
            goto cleanup;
        }
        mpz_sub_ui(index, shard->current, 1);
        if (mpz_lt(index, xsend.max_index)) {
            shard->state.hosts_allowlisted++;
            shard->iterations++;
            blocklist_lookup_index_for_ipvx_port_index(ipvx, index);
            goto cleanup;
        }
        shard->state.hosts_blocklisted++;
    }

cleanup:
    mpz_clear(index);
}

static void _shard_roll_to_valid(mpz_t ipvx, shard_t *shard) {
    if (mpz_le(shard->current, xsend.max_index)) {
        mpz_set(ipvx, shard->current);
        return;
    }
    _shard_get_next_ipvx_port_index(ipvx, shard);
}

void shard_init(shard_t *shard, uint16_t shard_idx, uint16_t num_shards,
                uint8_t thread_idx, uint8_t num_threads,
                uint64_t max_total_targets, uint64_t max_total_packets,
                uint64_t list_of_ips_count, const cycle_t *cycle,
                shard_complete_cb cb, void *args) {
    // Start out by figuring out how many shards we have. A single shard of
    // XMap (set with --shards=N, --shard=n) may have several subshards, if
    // XMap is being ran multithreaded (set with --sender-threads=T).
    //
    // Total number of subshards is S = N*T. Subshard ID's range from [0, N*T).
    log_debug("shard", "shard_init: shard-%d", thread_idx);

    assert(num_shards > 0);
    assert(num_threads > 0);
    assert(shard_idx < num_shards);
    assert(thread_idx < num_threads);
    uint32_t num_subshards = (uint32_t) num_shards * (uint32_t) num_threads;
    mpz_t    num_elts;
    mpz_init_set(num_elts, cycle->order);
    assert(mpz_gt_ui(num_elts, num_subshards)); // in case to many shards
    assert(!max_total_targets || max_total_targets >= num_subshards);
    assert(!max_total_packets || max_total_packets >= num_subshards);
    assert(!list_of_ips_count || list_of_ips_count >= num_subshards);

    // This instance of XMap will run T subshards, with one subshard per
    // thread. This composes a single shard, as specified by the command
    // line flag --shard=n.  E.g. to run shard with index n, we must run
    // subshards with indicies the range [n*T, (n+1)*T].
    //
    // We can calculate our subshard index i = n*T + t.
    uint32_t sub_idx = shard_idx * num_threads + thread_idx;

    // Given i, we want to calculate the start of subshard i. Subshards
    // define ranges over exponents of g. They range from [0, Q-1), where Q
    // is the number of elements in (the order of) the group generated by g.
    //
    // Let e_b = floor(Q / S) * i
    mpz_t exponent_begin;
    mpz_init(exponent_begin);
    mpz_fdiv_q_ui(exponent_begin, num_elts, num_subshards);
    mpz_mul_ui(exponent_begin, exponent_begin, sub_idx);

    // The stopping exponent is the first element of the next shard.
    //
    //    e_e = floor(Q / S) * ((i + 1) % S);
    mpz_t exponent_end;
    mpz_init(exponent_end);
    mpz_fdiv_q_ui(exponent_end, num_elts, num_subshards);
    mpz_mul_ui(exponent_end, exponent_end, (sub_idx + 1) % num_subshards);

    // We actually offset the begin and end of each cycle. Given an offset
    // k, shift each exponent by k modulo Q.
    mpz_add(exponent_begin, exponent_begin, cycle->offset);
    mpz_mod(exponent_begin, exponent_begin, num_elts);
    mpz_add(exponent_end, exponent_end, cycle->offset);
    mpz_mod(exponent_end, exponent_end, num_elts);

    // Calculate the first and last points of the shard as powers of g
    // modulo p. g^i mod p (i: index in a range)
    mpz_t start, stop;
    mpz_init(start);
    mpz_init(stop);
    mpz_powm(start, cycle->generator, exponent_begin, cycle->group->prime);
    mpz_powm(stop, cycle->generator, exponent_end, cycle->group->prime);

    // Pull the result out
    mpz_init_set(shard->params.first, start);
    mpz_init_set(shard->params.last, stop);
    mpz_init_set(shard->params.factor, cycle->generator);
    mpz_init_set(shard->params.modulus, cycle->group->prime);

    // Set the shard at the beginning.
    mpz_init_set(shard->current, shard->params.first);

    // Set the (thread) id
    shard->thread_id = thread_idx;

    // Set max_targets if applicable
    if (max_total_targets > 0) {
        uint64_t max_targets_this_shard = max_total_targets / num_subshards;
        if (sub_idx < (max_total_targets % num_subshards))
            ++max_targets_this_shard;
        shard->state.max_hosts = max_targets_this_shard;
    }

    // Set max_packets if applicable
    if (max_total_packets > 0) {
        uint64_t max_packets_this_shard = max_total_packets / num_subshards;
        if (sub_idx < (max_total_packets % num_subshards))
            ++max_packets_this_shard;
        shard->state.max_packets = max_packets_this_shard;
    }

    // initializing ip_target_file
    shard->ip_target_file_params.fp = NULL;
    if (list_of_ips_count) {
        uint64_t ip_target_this_shard      = list_of_ips_count / num_subshards;
        shard->ip_target_file_params.first = sub_idx * ip_target_this_shard;
        if (sub_idx < (list_of_ips_count % num_subshards)) {
            ++ip_target_this_shard;
            shard->ip_target_file_params.first += sub_idx;
        } else {
            shard->ip_target_file_params.first +=
                (list_of_ips_count % num_subshards);
        }
        shard->ip_target_file_params.current =
            shard->ip_target_file_params.first;
        shard->ip_target_file_params.total = ip_target_this_shard;
        shard->ip_target_file_params.last = shard->ip_target_file_params.first +
                                            shard->ip_target_file_params.total;
        shard->ip_target_file_params.fp =
            fopen(xconf.list_of_ips_filename, "r");
        if (shard->ip_target_file_params.fp == NULL) {
            log_fatal("shard", "open ip target file for shard failed");
        }
        shard->ip_target_file_params.port_current  = 0;
        shard->ip_target_file_params.port_total    = xconf.target_port_num;
        shard->ip_target_file_params.index_current = 0;
        shard->ip_target_file_params.index_total   = xconf.target_index_num;
    }

    // Set the callbacks
    shard->completeCb = cb;
    shard->args       = args;

    // If the beginning of a shard isn't pointing to a valid index in the
    // blocklist, find the first element that is.
    // num_elts just take a position
    _shard_roll_to_valid(num_elts, shard);

    // initializing other
    mpz_init_set(shard->state.first_scanned, shard->current);

    // Clear everything
    mpz_clear(num_elts);
    mpz_clear(exponent_begin);
    mpz_clear(exponent_end);
    mpz_clear(start);
    mpz_clear(stop);

    log_debug("shard", "shard_init: shard-%d completed", thread_idx);
}

void shard_get_current_ip_prefix_port_index(void *prefix, shard_t *shard,
                                            port_h_t  *port_f,
                                            index_h_t *index_f) {
    mpz_t ipvx_port_index, ipvx, port_m, index_m;
    mpz_init(ipvx_port_index);
    mpz_init(ipvx);
    mpz_init(port_m);
    mpz_init(index_m);

    _shard_get_current_ipvx_port_index(ipvx_port_index, shard);
    mpz_mod_2exp(port_m, ipvx_port_index, xconf.max_port_index_len);
    mpz_fdiv_q_2exp(port_m, port_m, xconf.target_index_bits);
    mpz_mod_2exp(index_m, ipvx_port_index, xconf.target_index_bits);
    port_h_t  port  = (port_h_t) mpz_get_ui(port_m);
    index_h_t index = (index_h_t) mpz_get_ui(index_m);

    while (
        mpz_ne_ui(ipvx_port_index, PMAP_SHARD_DONE) &&
        ((xconf.target_port_num && port >= xconf.target_port_num) ||
         (xconf.target_index_num && (int) index >= xconf.target_index_num))) {
        _shard_get_next_ipvx_port_index(ipvx_port_index, shard);
        mpz_mod_2exp(port_m, ipvx_port_index, xconf.max_port_index_len);
        mpz_fdiv_q_2exp(port_m, port_m, xconf.target_index_bits);
        mpz_mod_2exp(index_m, ipvx_port_index, xconf.target_index_bits);
        port  = (port_h_t) mpz_get_ui(port_m);
        index = (index_h_t) mpz_get_ui(index_m);
    }

    mpz_fdiv_q_2exp(ipvx, ipvx_port_index, xconf.max_port_index_len);
    memset(prefix, 0, xconf.ipv46_bytes);
    mpz_to_uint8s_bits(ipvx, prefix, xconf.max_probe_len);

    mpz_clear(ipvx_port_index);
    mpz_clear(ipvx);
    mpz_clear(port_m);
    mpz_clear(index_m);

    *port_f  = xconf.target_port_list[port];
    *index_f = index;
}

void shard_get_next_ip_prefix_port_index(void *prefix, shard_t *shard,
                                         port_h_t *port_f, index_h_t *index_f) {
    mpz_t ipvx_port_index, ipvx, port_m, index_m;
    mpz_init(ipvx_port_index);
    mpz_init(ipvx);
    mpz_init(port_m);
    mpz_init(index_m);

    _shard_get_next_ipvx_port_index(ipvx_port_index, shard);
    mpz_mod_2exp(port_m, ipvx_port_index, xconf.max_port_index_len);
    mpz_fdiv_q_2exp(port_m, port_m, xconf.target_index_bits);
    mpz_mod_2exp(index_m, ipvx_port_index, xconf.target_index_bits);
    port_h_t  port  = (port_h_t) mpz_get_ui(port_m);
    index_h_t index = (index_h_t) mpz_get_ui(index_m);

    while (
        mpz_ne_ui(ipvx_port_index, PMAP_SHARD_DONE) &&
        ((xconf.target_port_num && port >= xconf.target_port_num) ||
         (xconf.target_index_num && (int) index >= xconf.target_index_num))) {
        _shard_get_next_ipvx_port_index(ipvx_port_index, shard);
        mpz_mod_2exp(port_m, ipvx_port_index, xconf.max_port_index_len);
        mpz_fdiv_q_2exp(port_m, port_m, xconf.target_index_bits);
        mpz_mod_2exp(index_m, ipvx_port_index, xconf.target_index_bits);
        port  = (port_h_t) mpz_get_ui(port_m);
        index = (index_h_t) mpz_get_ui(index_m);
    }

    mpz_fdiv_q_2exp(ipvx, ipvx_port_index, xconf.max_port_index_len);
    memset(prefix, 0, xconf.ipv46_bytes);
    mpz_to_uint8s_bits(ipvx, prefix, xconf.max_probe_len);

    mpz_clear(ipvx_port_index);
    mpz_clear(ipvx);
    mpz_clear(port_m);
    mpz_clear(index_m);

    *port_f  = xconf.target_port_list[port];
    *index_f = index;
}

void shard_free(shard_t *shard) {
    mpz_clear(shard->state.first_scanned);
    mpz_clear(shard->params.first);
    mpz_clear(shard->params.last);
    mpz_clear(shard->params.factor);
    mpz_clear(shard->params.modulus);
    if (shard->ip_target_file_params.fp)
        fclose(shard->ip_target_file_params.fp);
    mpz_clear(shard->current);
}