#include "Firewall.h"

ESPFirewall::ESPFirewall(int port)
{
    this->setup_eeprom();
    log_i("Starting Firewall-API on %i", port);
    this->firewall_api = new AsyncWebServer(port);
    this->setup_routing();
}

String ESPFirewall::protocol_to_string(firewall_protocol_t &protocol)
{
    switch (protocol)
    {
    case FW_TCP:
        return "TCP";
    case FW_UDP:
        return "UDP";
    default:
        return "ALL";
    }
}

firewall_protocol_t ESPFirewall::string_to_protocol(String &protocol)
{
    if (protocol.equals("TCP"))
        return FW_TCP;
    else if (protocol.equals("UDP"))
        return FW_UDP;
    else
        return FW_ALL;
}

String ESPFirewall::target_to_string(firewall_target_t &target)
{
    switch (target)
    {
    case FW_REJECT:
        return "REJECT";
    case FW_DROP:
        return "DROP";
    default:
        return "ACCEPT";
    }
}

firewall_target_t ESPFirewall::string_to_target(String &target)
{
    if (target.equals("REJECT"))
        return FW_REJECT;
    else if (target.equals("DROP"))
        return FW_DROP;
    else
        return FW_ACCEPT;
}

void ESPFirewall::setup_eeprom()
{
    EEPROM.begin(this->eeprom_size);
    this->amount_of_rules = EEPROM.read(this->eeprom_settings_head);
    uint8_t security_number = EEPROM.read(this->eeprom_settings_head + 1);
    log_i("Amount of existing Rules %i", this->amount_of_rules);
    if (this->amount_of_rules > 50 || security_number != this->security_number)
    {
        this->amount_of_rules = 0;
        EEPROM.write(this->eeprom_settings_head, this->amount_of_rules);
        EEPROM.write(this->eeprom_settings_head + 1, this->security_number);
        EEPROM.commit();
    }
    log_i("Amount of existing Rules %i", this->amount_of_rules);
    this->eeprom_read_firewall_rules();
}

void ESPFirewall::eeprom_write_firewall_rule(firewall_rule_t *rule_ptr)
{
    EEPROM.write(this->eeprom_settings_head, this->amount_of_rules);
    EEPROM.writeString(this->eeprom_rules_head, rule_ptr->source);
    this->eeprom_rules_head += IP4ADDR_STRLEN_MAX;
    EEPROM.writeString(this->eeprom_rules_head, rule_ptr->destination);
    this->eeprom_rules_head += IP4ADDR_STRLEN_MAX;
    EEPROM.write(this->eeprom_rules_head, rule_ptr->protocol);
    this->eeprom_rules_head += sizeof(firewall_protocol_t);
    EEPROM.write(this->eeprom_rules_head, rule_ptr->target);
    this->eeprom_rules_head += sizeof(firewall_target_t);
    EEPROM.commit();
}

void ESPFirewall::eeprom_write_firewall_rules()
{
    this->eeprom_rules_head = eeprom_start_firewall_rules;
    firewall_rule_t *rule_ptr = this->head;
    while (rule_ptr != NULL)
    {
        this->eeprom_write_firewall_rule(rule_ptr);
        rule_ptr = rule_ptr->next;
    }
}

void ESPFirewall::eeprom_read_firewall_rule(uint8_t &eeprom_address, uint8_t &rule_nr)
{
    firewall_rule_t *rule_ptr = (firewall_rule_t *)malloc(sizeof(firewall_rule_t));
    rule_ptr->key = rule_nr;
    strcpy(rule_ptr->source, EEPROM.readString(eeprom_address).c_str());
    eeprom_address += IP4ADDR_STRLEN_MAX;
    strcpy(rule_ptr->destination, EEPROM.readString(eeprom_address).c_str());
    eeprom_address += IP4ADDR_STRLEN_MAX;
    rule_ptr->protocol = static_cast<firewall_protocol_t>(EEPROM.read(eeprom_address));
    eeprom_address += sizeof(firewall_protocol_t);
    rule_ptr->target = static_cast<firewall_target_t>(EEPROM.read(eeprom_address));
    eeprom_address += sizeof(firewall_target_t);
    add_rule_to_firewall(rule_ptr);
    log_i("%s, %s, %s, %s",
          rule_ptr->source,
          rule_ptr->destination,
          protocol_to_string(rule_ptr->protocol),
          target_to_string(rule_ptr->target));
}

void ESPFirewall::eeprom_read_firewall_rules()
{
    uint8_t eeprom_address = eeprom_start_firewall_rules;
    for (uint8_t i = 1; i <= this->amount_of_rules; i++)
    {
        eeprom_read_firewall_rule(eeprom_address, i);
    }
}

void ESPFirewall::add_rule_to_firewall(firewall_rule_t *rule_ptr)
{
    firewall_rule_t *temp;
    if (this->head == NULL)
    {
        this->head = rule_ptr;
        rule_ptr->next = NULL;
        return;
    }
    temp = this->head;
    while (temp->next != NULL)
    {
        temp = temp->next;
    }
    temp->next = rule_ptr;
    rule_ptr->next = NULL;
    return;
}

firewall_rule_t *ESPFirewall::get_rule_from_firewall(uint8_t key)
{
    firewall_rule_t *rule_ptr = this->head;
    if (this->head == NULL)
    {
        return NULL;
    }
    while (rule_ptr->key != key)
    {
        if (rule_ptr->next == NULL)
        {
            return NULL;
        }
        else
        {
            rule_ptr = rule_ptr->next;
        }
    }
    return rule_ptr;
}

bool ESPFirewall::delete_rule_from_firewall(uint8_t key)
{
    if (this->head == NULL)
    {
        return false;
    }
    firewall_rule_t *current_rule_ptr = this->head;
    firewall_rule_t *previous_rule_ptr = NULL;
    firewall_rule_t *temp = NULL;
    while (current_rule_ptr->key != key)
    {
        if (current_rule_ptr->next == NULL)
        {
            return false;
        }
        else
        {
            previous_rule_ptr = current_rule_ptr;
            current_rule_ptr = current_rule_ptr->next;
        }
    }
    if (current_rule_ptr == this->head)
    {
        this->head = head->next;
        temp = this->head;
    }
    else
    {
        previous_rule_ptr->next = current_rule_ptr->next;
        temp = previous_rule_ptr->next;
    }
    while (temp != NULL)
    {
        temp->key--;
        temp = temp->next;
    }
    free(current_rule_ptr);
    this->amount_of_rules--;
    this->eeprom_write_firewall_rules();
    return true;
}

void ESPFirewall::setup_routing()
{
    firewall_api->on("^\\/api/v1/firewall\\/([0-9]+)$", HTTP_GET, std::bind(&ESPFirewall::get_firewall_rule_handler, this, std::placeholders::_1));
    firewall_api->on("/api/v1/firewall", HTTP_GET, std::bind(&ESPFirewall::get_firewall_rules_handler, this, std::placeholders::_1));
    firewall_api->on("/api/v1/firewall", HTTP_POST, std::bind(&ESPFirewall::post_firewall_handler, this, std::placeholders::_1));
    firewall_api->on("^\\/api/v1/firewall\\/([0-9]+)$", HTTP_DELETE, std::bind(&ESPFirewall::delete_firewall_handler, this, std::placeholders::_1));

    firewall_api->on("/api/v1/device/restart", HTTP_GET, std::bind(&ESPFirewall::restart_device_handler, this, std::placeholders::_1));
    firewall_api->onNotFound(std::bind(&ESPFirewall::not_found, this, std::placeholders::_1));
    this->firewall_api->begin();
}

void ESPFirewall::json_message_response(AsyncWebServerRequest *request, String message, int response_code)
{
    DynamicJsonDocument json(256);
    String response;
    json["message"] = message;
    serializeJson(json, response);
    request->send(response_code, "application/json", response);
}

String ESPFirewall::construct_json_firewall_rule(firewall_rule_t *rule_ptr)
{
    StaticJsonDocument<192> doc;
    doc["key"] = rule_ptr->key;
    doc["source"] = rule_ptr->source;
    doc["destination"] = rule_ptr->destination;
    doc["protocol"] = protocol_to_string(rule_ptr->protocol);
    doc["target"] = target_to_string(rule_ptr->target);
    String response;
    serializeJson(doc, response);
    return response;
}

String ESPFirewall::construct_json_firewall()
{
    firewall_rule_t *rule_ptr = this->head;
    // Size for max 12 Rules
    StaticJsonDocument<2048> doc;
    String response;
    doc["amount_of_rules"] = this->amount_of_rules;
    JsonArray rules = doc.createNestedArray("rules");
    while (rule_ptr != NULL)
    {
        JsonObject rule = rules.createNestedObject();
        rule["key"] = rule_ptr->key;
        rule["source"] = rule_ptr->source;
        rule["destination"] = rule_ptr->destination;
        rule["protocol"] = protocol_to_string(rule_ptr->protocol);
        rule["target"] = target_to_string(rule_ptr->target);
        rule_ptr = rule_ptr->next;
    }
    serializeJson(doc, response);
    return response;
}

void ESPFirewall::not_found(AsyncWebServerRequest *request)
{
    json_message_response(request, "not found", 404);
}

void ESPFirewall::restart_device_handler(AsyncWebServerRequest *request)
{
    json_message_response(request, "restarting device in 2 sec", 200);
    sleep(2000);
    esp_restart();
}

void ESPFirewall::get_firewall_rule_handler(AsyncWebServerRequest *request)
{
    int rule_number = request->pathArg(0).toInt();
    firewall_rule_t *rule_ptr = get_rule_from_firewall(rule_number);
    if (rule_ptr == NULL)
    {
        json_message_response(request, "rule not found", 404);
    }
    else
    {
        request->send(200, "application/json", construct_json_firewall_rule(rule_ptr));
    }
}

void ESPFirewall::get_firewall_rules_handler(AsyncWebServerRequest *request)
{
    String response = construct_json_firewall();
    request->send(200, "application/json", response);
}

bool ESPFirewall::request_has_firewall_parameter(AsyncWebServerRequest *request)
{
    return request->hasParam("source") || request->hasParam("destination") || request->hasParam("protocol") || request->hasParam("target");
}

void ESPFirewall::post_firewall_handler(AsyncWebServerRequest *request)
{
    if (request_has_firewall_parameter(request))
    {
        firewall_rule_t *rule_ptr = (firewall_rule_t *)malloc(sizeof(firewall_rule_t));
        rule_ptr->key = ++amount_of_rules;

        // carefully copying c-string that is shorter then the destination char-array length
        String source = request->getParam("source")->value();
        strcpy(rule_ptr->source, source.length() <= IP4ADDR_STRLEN_MAX ? source.c_str() : "");
        String destination = request->getParam("destination")->value();
        strcpy(rule_ptr->destination, destination.length() <= IP4ADDR_STRLEN_MAX ? destination.c_str() : "");

        String protocol = request->getParam("protocol")->value();
        rule_ptr->protocol = string_to_protocol(protocol);
        String target = request->getParam("target")->value();
        rule_ptr->target = string_to_target(target);

        add_rule_to_firewall(rule_ptr);
        eeprom_write_firewall_rule(rule_ptr);
        request->send(200, "application/json", construct_json_firewall_rule(rule_ptr));
    }
    else
    {
        json_message_response(request, "not enough parameter", 200);
    }
}

void ESPFirewall::delete_firewall_handler(AsyncWebServerRequest *request)
{
    int rule_number = request->pathArg(0).toInt();
    if (delete_rule_from_firewall(rule_number))
    {
        json_message_response(request, "firewall rule deleted", 200);
    }
    json_message_response(request, "cannot delete firewall rule", 500);
}