
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <arpa/inet.h>
#include <sys/socket.h>

#include "log.h"
#include "conf.h"
#include "ext-fwd.h"
#include "natpmp.h"


enum {
    NATPMP_STATE_INIT,
    NATPMP_STATE_REQUEST_GATEWAY,
    NATPMP_STATE_RECEIVE_GATEWAY,
    NATPMP_STATE_REQUEST_PORTMAPPING,
    NATPMP_STATE_RECEIVE_PORTMAPPING,
    NATPMP_STATE_ERROR
};

#ifdef DEBUG
const char* natpmp_statestr(int state)
{
    switch (state) {
    case NATPMP_STATE_INIT:
        return "NATPMP_STATE_INIT";
    case NATPMP_STATE_REQUEST_GATEWAY:
        return "NATPMP_STATE_REQUEST_GATEWAY";
    case NATPMP_STATE_RECEIVE_GATEWAY:
        return "NATPMP_STATE_RECEIVE_GATEWAY";
    case NATPMP_STATE_REQUEST_PORTMAPPING:
        return "NATPMP_STATE_REQUEST_PORTMAPPING";
    case NATPMP_STATE_RECEIVE_PORTMAPPING:
        return "NATPMP_STATE_RECEIVE_PORTMAPPING";
    case NATPMP_STATE_ERROR:
        return "NATPMP_STATE_ERROR";
    default:
        return "<unknown>";
    }
}
#endif

void natpmp_init(struct natpmp_handle_t **handle)
{
    // Initialize data structure
    struct natpmp_handle_t *m = calloc(1, sizeof(struct natpmp_handle_t));
    m->state = 0;
    m->retry = 0;
    m->natpmp.s = -1;
    *handle = m;
}

void natpmp_uninit(struct natpmp_handle_t **handle)
{
    // Remove all port mapping associated with this host
    struct natpmp_handle_t *m = *handle;
    sendnewportmappingrequest(&m->natpmp, NATPMP_PROTOCOL_TCP, 0, 0, 0);
    sendnewportmappingrequest(&m->natpmp, NATPMP_PROTOCOL_UDP, 0, 0, 0);
    closenatpmp(&m->natpmp);
    free(*handle);
    *handle = NULL;
}

int natpmp_handler(struct natpmp_handle_t *handle, uint16_t port, time_t lifespan, time_t now)
{
    natpmpresp_t response;

    // Retry later if we want to wait longer
    if (handle->retry > now) {
        return PF_RETRY;
    }

#ifdef DEBUG
    log_debug("NAT-PMP: Handle port: %hu, lifespan: %ld, state: %s", port, lifespan, natpmp_statestr(handle->state));
#endif

    // Initialize data structure / socket
    if (handle->state == NATPMP_STATE_INIT) {
        int rc = initnatpmp(&handle->natpmp, 0, 0);
        if (rc >= 0) {
            handle->state = NATPMP_STATE_REQUEST_GATEWAY;
            return PF_RETRY;
        } else {
            log_debug("NAT-PMP: Method initnatpmp returned %d.", rc);
            goto error;
        }
    }

    // Request gateway address
    if (handle->state == NATPMP_STATE_REQUEST_GATEWAY) {
        int rc = sendpublicaddressrequest(&handle->natpmp);
        if (rc >= 0) {
            handle->retry = now + 8;
            handle->state = NATPMP_STATE_RECEIVE_GATEWAY;
            return PF_RETRY;
        } else {
            log_debug("NAT-PMP: Method sendpublicaddressrequest returned %d.", rc);
            goto error;
        }
    }

    // Read public gateway address
    if (handle->state == NATPMP_STATE_RECEIVE_GATEWAY) {
        int rc = readnatpmpresponseorretry(&handle->natpmp, &response);
        
        if (rc >= 0) {
            char str[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &response.pnu.publicaddress.addr, str, sizeof (str));
            log_info("NAT-PMP: Found public address \"%s\".", str);
            handle->state = NATPMP_STATE_REQUEST_PORTMAPPING;
            return PF_RETRY;
        } else if (rc == NATPMP_TRYAGAIN) {
            handle->retry = now + (10 * 60);
            handle->state = NATPMP_STATE_REQUEST_GATEWAY;
            return PF_RETRY;
        } else {
            log_debug("NAT-PMP: Method readnatpmpresponseorretry returned %d.", rc);
            goto error;
        }
    }

    // Add/Remove port mappings
    if (handle->state == NATPMP_STATE_REQUEST_PORTMAPPING) {
        int rc_udp = sendnewportmappingrequest(&handle->natpmp, NATPMP_PROTOCOL_UDP, port, port, lifespan);
        int rc_tcp = sendnewportmappingrequest(&handle->natpmp, NATPMP_PROTOCOL_TCP, port, port, lifespan);

        if (rc_udp >= 0 && rc_tcp >= 0) {
            handle->retry = now + 2;
            handle->state = NATPMP_STATE_RECEIVE_PORTMAPPING;
            return PF_RETRY;
        } else {

            log_debug("NAT-PMP: Method sendnewportmappingrequest returned an error: %s",
                strnatpmperr((rc_udp >= 0) ? rc_tcp : rc_udp));
            goto error;
        }
    }

    // Check port mapping
    if (handle->state == NATPMP_STATE_RECEIVE_PORTMAPPING) {
        int rc = readnatpmpresponseorretry(&handle->natpmp, &response);
        if (rc >= 0) {
            int private_port = response.pnu.newportmapping.privateport;
            int public_port = response.pnu.newportmapping.mappedpublicport;
            time_t lifetime = response.pnu.newportmapping.lifetime;

            if (lifetime > 0) {
                log_info("NAT-PMP: Port forwarding added for port %d (to private port %d) for %ld seconds.", public_port, private_port, lifetime);
                handle->state = NATPMP_STATE_REQUEST_PORTMAPPING;
                return PF_DONE;
            } else {
                log_debug("NAT-PMP: Port forwarding removed for public port %d (to private port %d) for %ld seconds.", public_port, private_port, lifetime);
                handle->state = NATPMP_STATE_REQUEST_PORTMAPPING;
                return PF_DONE;
            }
        } else if (rc == NATPMP_TRYAGAIN) {
            handle->state = NATPMP_STATE_RECEIVE_PORTMAPPING;
            return PF_RETRY;
        } else {
            log_debug("NAT-PMP: Port forwarding failed for port %d.", port);
            goto error;
        }
    }

    error:;
    handle->retry = now + 60;
    handle->state = NATPMP_STATE_ERROR;
    return PF_ERROR;
}
