#include <nginx.h>
#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>

#include <openssl/evp.h>
#include <openssl/aes.h>

static const u_int32_t BLOCK_SIZE = 16;

static ngx_int_t ngx_http_aes_filter_init(ngx_conf_t *cf);
static void * ngx_http_aes_create_conf(ngx_conf_t *cf);
static char * ngx_http_aes_merge_conf(ngx_conf_t *cf, void *parent, void *child);

typedef struct {
    ngx_uint_t    bufs;

    u_int16_t     range_read;

    ngx_chain_t*  free;
    ngx_chain_t*  busy;
    ngx_chain_t*  out;
    ngx_chain_t*  in;

    unsigned char ivec[16];
    unsigned int  num;
    unsigned char ecount[16];
    AES_KEY       aes_key;
} ngx_http_aes_ctx_t;

// get http_range data from the http_range module
extern ngx_module_t  ngx_http_range_body_filter_module;
typedef struct {
    off_t        start;
    off_t        end;
    ngx_str_t    content_range;
} ngx_http_range_t;
typedef struct {
    off_t        offset;
    ngx_str_t    boundary_header;
    ngx_array_t  ranges;
} ngx_http_range_filter_ctx_t;

typedef struct {
    ngx_flag_t enable;
    ngx_bufs_t bufs;
} ngx_http_aes_conf_t;

static ngx_command_t  ngx_http_aes_filter_commands[] = {
    { ngx_string("aes"),
        NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_FLAG,
        ngx_conf_set_flag_slot,
        NGX_HTTP_LOC_CONF_OFFSET,
        offsetof(ngx_http_aes_conf_t, enable),
        NULL },
    { ngx_string("aes_buffers"),
          NGX_HTTP_MAIN_CONF|NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_TAKE2,
          ngx_conf_set_bufs_slot,
          NGX_HTTP_LOC_CONF_OFFSET,
          offsetof(ngx_http_aes_conf_t, bufs),
          NULL },
    ngx_null_command
};

static ngx_http_module_t ngx_http_aes_module_ctx = {
    NULL,                                 /* preconfiguration */
    ngx_http_aes_filter_init,             /* postconfiguration */
    NULL,                                 /* create main configuration */
    NULL,                                 /* init main configuration */
    NULL,                                 /* create server configuration */
    NULL,                                 /* merge server configuration */
    ngx_http_aes_create_conf,            /* create location configuration */
    ngx_http_aes_merge_conf              /* merge location configuration */
};

ngx_module_t ngx_http_aes_module = {
    NGX_MODULE_V1,
    &ngx_http_aes_module_ctx, /* module context */
    ngx_http_aes_filter_commands,          /* module directives */
    NGX_HTTP_MODULE,                 /* module type */
    NULL,                            /* init master */
    NULL,                            /* init module */
    NULL,                            /* init process */
    NULL,                            /* init thread */
    NULL,                            /* exit thread */
    NULL,                            /* exit process */
    NULL,                            /* exit master */
    NGX_MODULE_V1_PADDING
};

static ngx_http_output_header_filter_pt  ngx_http_next_header_filter;
static ngx_http_output_body_filter_pt  ngx_http_next_body_filter;

static ngx_table_elt_t *
search_headers(ngx_http_upstream_headers_in_t *headers, u_char *name, size_t len) {
    ngx_list_part_t *part;
    ngx_table_elt_t *h;
    ngx_uint_t i;

    part = &headers->headers.part;
    h = part->elts;

    for (i = 0; /* void */; i++) {
        if (i >= part->nelts) {
            if (part->next == NULL) {
                break;
            }
            part = part->next;
            h = part->elts;
            i = 0;
        }
        if (len != h[i].key.len || ngx_strcasecmp(name, h[i].key.data) != 0) {
            continue;
        }
        return &h[i];
    }
    return NULL;
}

ngx_http_aes_ctx_t*
init_ctx (ngx_http_request_t* r)
{
    ngx_http_aes_ctx_t* ctx = ngx_pcalloc (r->pool, sizeof(ngx_http_aes_ctx_t));
    if (ctx == NULL) {
        return ctx;
    }

    ctx->bufs = 0;
    ctx->range_read = 0;

    ctx->free = NULL;
    ctx->busy = NULL;
    ctx->out = NULL;

//    ctx->ivec = ;
    ctx->num = 0;
//    ctx->ecount =;
//    ctx->aes_key =;

    return ctx;
}

static ngx_int_t
ngx_http_aes_header_filter(ngx_http_request_t *r)
{
    ngx_http_aes_conf_t* conf = ngx_http_get_module_loc_conf(r, ngx_http_aes_module);
    if (!conf->enable) {
        return ngx_http_next_header_filter(r);
    }

    ngx_table_elt_t* header_val = search_headers(&r->upstream->headers_in, "x-accel-aes", strlen("x-accel-aes"));
    if (header_val == NULL) {
        return ngx_http_next_header_filter(r);
    }

    ngx_http_aes_ctx_t* ctx = init_ctx(r);
    if (ctx == NULL) {
        return NGX_ERROR;
    }
    ngx_http_set_ctx(r, ctx, ngx_http_aes_module);

    ngx_str_t raw_key = header_val->value;

    u_char* key = ngx_pcalloc(r->pool, BLOCK_SIZE);
    PKCS5_PBKDF2_HMAC(raw_key.data, raw_key.len,
                      "8&AqGhV#nz sXVG1N+]e@>c +=c#D,wx", 32, 2,
                      EVP_sha512(),
                      BLOCK_SIZE, key);

    u_char* iv = ngx_pcalloc(r->pool, BLOCK_SIZE);
    PKCS5_PBKDF2_HMAC(raw_key.data, raw_key.len,
                      "Vr  '.J9ipG(]=Z!=YPy\\i5l3G(l8JG^3", 33, 2,
                      EVP_sha512(),
                      BLOCK_SIZE, iv);

    ctx->num = 0;
    memset(ctx->ecount, 0, BLOCK_SIZE);
    memcpy(ctx->ivec, iv, BLOCK_SIZE);

    int result = AES_set_encrypt_key(key, 128, &ctx->aes_key);
    if (result < 0)
    {
        ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "Error trying to set the AES key.");
        return NGX_ERROR;
    }

    return ngx_http_next_header_filter(r);
}

ngx_int_t
get_chain_elm(ngx_http_request_t* r, ngx_http_aes_ctx_t* ctx, size_t size, ngx_chain_t** new_elm)
{
    ngx_chain_t* chain = ngx_chain_get_free_buf(r->pool, &ctx->free);
    if (chain == NULL)
    {
        return NGX_ERROR;
    }

    *new_elm = chain;
    if (chain->buf->start == NULL)
    {
        ngx_http_aes_conf_t* conf = ngx_http_get_module_loc_conf(r, ngx_http_aes_module);
        if (ctx->bufs < conf->bufs.num) {
            chain->buf->start = ngx_palloc(r->pool, size);
            if (chain->buf->start == NULL)
            {
                return NGX_ERROR;
            }
            chain->buf->pos = chain->buf->start;
            chain->buf->last = chain->buf->start;
            chain->buf->end = chain->buf->start + size;
            chain->buf->memory = 1;
            //chain->buf->recycled = 1;
//            chain->buf->flush = 1;
            chain->buf->tag = (ngx_buf_tag_t) &ngx_http_aes_module;
            ctx->bufs++;
        } else {
            return NGX_DECLINED;
        }
    }
    return NGX_OK;
}

static void ctr128_inc(unsigned char *counter)
{
    uint32_t n = 16;
    unsigned char c;

    do {
        --n;
        c = counter[n];
        ++c;
        counter[n] = c;
        if (c)
            return;
    } while (n);
}

static ngx_int_t
ngx_http_aes_next_filter_update_chains(ngx_http_request_t *r, ngx_http_aes_ctx_t* ctx, ngx_chain_t *in)
{
    ngx_int_t rc, rc2;
    rc = ngx_http_next_body_filter(r, ctx->out);
    ngx_chain_update_chains(r->pool, &ctx->free, &ctx->busy, &ctx->out,
                            (ngx_buf_tag_t) &ngx_http_aes_module);

    if (rc == NGX_AGAIN && in != NULL && in->buf != NULL && ngx_buf_size(in->buf) > 0) {

        // save the current in buffer for further use
        ctx->in = NULL;
        rc2 = ngx_chain_add_copy(r->pool, &ctx->in, in);

        if (rc2 != NGX_OK) {
            ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "Cannot allocate new chain");
            return NGX_ERROR;
        }
    }

    return rc;
}

static ngx_int_t logg(ngx_http_request_t *r, int line, ngx_int_t a) {
//    ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "return %i, line=%i", a, line);
    return a;
}

static ngx_int_t ngx_http_aes_body_filter(ngx_http_request_t *r, ngx_chain_t *in)
{
    ngx_int_t rc;
    ngx_http_aes_ctx_t* ctx = ngx_http_get_module_ctx(r, ngx_http_aes_module);
    if (ctx == NULL) {
        return logg(r, __LINE__, ngx_http_next_body_filter(r, in));
    }

    // If we do not start at the beginning of the file
    // we need to initialize the AES_crt_128 algorithm
    if (!ctx->range_read) {
        ctx->range_read = 1;
        ngx_http_range_filter_ctx_t* ctx_range = ngx_http_get_module_ctx(r, ngx_http_range_body_filter_module);
        if (ctx_range != NULL) {
            off_t start = 0;
            if (ctx_range->ranges.nelts == 1 && (start = ((ngx_http_range_t*)ctx_range->ranges.elts)->start) > 0) {
                ctx->num = start % BLOCK_SIZE;
                uint32_t n = start / BLOCK_SIZE + 1;

                // openssl uses num to find if it's the first iteration
                if (ctx->num == 0) {
                    n--;
                }
                while (n > 0) {
                    AES_encrypt(ctx->ivec, ctx->ecount, &ctx->aes_key);
                    ctr128_inc(ctx->ivec);
                    n--;
                }

            } else if (ctx_range->ranges.nelts > 1) {
                ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "Cannot aes decrypt multi-range request");
                return NGX_ERROR;
            }
        }
    }

    if (in == NULL) {
        rc = ngx_http_next_body_filter(r, in);

        if (rc != NGX_OK || ctx->in == NULL) {
            return logg(r, __LINE__, rc);
        }

        // 'in' was null, but we have some data not sent yet
        // ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "Uses old IN");
        in = ctx->in;
    }

    ngx_chain_t *cl_in = in;
    ngx_chain_t *curr_out = NULL;
    ngx_buf_t* input = NULL;
    ngx_http_aes_conf_t* conf = ngx_http_get_module_loc_conf(r, ngx_http_aes_module);
    ngx_int_t retry = 0;

    for ( ; ; ) {
        size_t size = ngx_buf_size(cl_in->buf);

        // Has long has this cl_in is empty, try the next one.
        if (size == 0) {
            // Everything is finished
            if (cl_in->buf->last_buf) {

                // curr_out does not exist yet => we just got an empty 'in'
                if (curr_out == NULL) {
                    return ngx_http_next_body_filter(r, NULL);
                }

                curr_out->buf->last_buf = 1;
                curr_out->buf->last_in_chain = 1;
                curr_out->buf->flush = 1;
                return logg(r, __LINE__, ngx_http_aes_next_filter_update_chains(r, ctx, cl_in));
            }

            // Still some data waiting somewhere
            cl_in = cl_in->next;

            if (cl_in == NULL) {
                // No more buffer in this chain
                if (curr_out != NULL) {
                    curr_out->buf->last_in_chain = 1;
                    curr_out->buf->flush = 1;
                }
                return logg(r, __LINE__, ngx_http_aes_next_filter_update_chains(r, ctx, cl_in));
            }

            // RESTART with the new input buffer
            continue;
        }

        // if output buffer is full/non-existing => create new one
        if (ctx->out == NULL || curr_out == NULL || curr_out->buf->end == curr_out->buf->last) {
            ngx_chain_t* chain;
            rc = get_chain_elm(r, ctx, conf->bufs.size, &chain);
            if (rc == NGX_ERROR) {
                return NGX_ERROR;
            }

            // No buffer available ...
            if (rc == NGX_DECLINED) {

                // ... empty the ones we have ...
                if (ctx->out != NULL) {
                    ctx->out->buf->flush = 1;
                }
                if (curr_out != NULL) {
                    curr_out->buf->flush = 1;
                }
                rc = logg(r, __LINE__, ngx_http_aes_next_filter_update_chains(r, ctx, cl_in));
                if (rc == NGX_AGAIN && retry-- > 0) {
                    ngx_log_error(NGX_LOG_ERR, r->connection->log, 0, "Retry=%i", retry);
                    continue;
                }
                if (rc != NGX_OK) {
                    return rc;
                }

                // ... RETRY with more buffers available
                continue;
            }

            // update current output
            if (ctx->out == NULL) {
                ctx->out = chain;
                curr_out = chain;
            } else {
                curr_out->next = chain;
                curr_out = chain;
            }
        }

        // If input has been read entirely or is null
        if (input == NULL || input->pos == input->last) {
            if (!ngx_buf_in_memory(cl_in->buf))
            {
                input = ngx_calloc_buf(r->pool);
                if (input == NULL) {
                    return NGX_ERROR;
                }
                input->start = curr_out->buf->start;
                input->end = curr_out->buf->end;
                input->pos = curr_out->buf->last;
                input->last = curr_out->buf->last;
                input->memory = 1;
                input->tag = (ngx_buf_tag_t) &ngx_http_aes_module;

                size_t read_size = ngx_min(size, input->end - input->pos);
                ssize_t n = ngx_read_file(cl_in->buf->file,
                                          input->pos,
                                          read_size,
                                          cl_in->buf->file_pos);
                if (n == NGX_ERROR) {
                    return NGX_ERROR;
                }

                if (n != read_size) {
                    ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
                                       ngx_read_file_n " returned "
                                       "only %z bytes instead of %z, %z, %z",
                                       n, read_size, size, input->end - input->pos);
                }

                input->last = input->pos + n;
                input->last_buf = cl_in->buf->last_buf;
                input->last_in_chain = cl_in->buf->last_in_chain;
                cl_in->buf->file_pos += n;
            } else {
                input = cl_in->buf;
            }
        }

        size_t usable_size = ngx_min(size, (curr_out->buf->end - curr_out->buf->last));

        AES_ctr128_encrypt(input->pos,
                           curr_out->buf->last,
                           usable_size,
                           &ctx->aes_key,
                           ctx->ivec,
                           ctx->ecount,
                           &ctx->num);

        input->pos += usable_size;
        curr_out->buf->last += usable_size;
    }
}

static void *
ngx_http_aes_create_conf(ngx_conf_t *cf)
{
    ngx_http_aes_conf_t  *conf;

    conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_aes_conf_t));
    if (conf == NULL) {
        return NULL;
    }

    conf->enable = NGX_CONF_UNSET;
    return conf;
}

static char *
ngx_http_aes_merge_conf(ngx_conf_t *cf, void *parent, void *child)
{
    ngx_http_aes_conf_t *prev = parent;
    ngx_http_aes_conf_t *conf = child;

    ngx_conf_merge_value(conf->enable, prev->enable, 0);

    ngx_conf_merge_bufs_value(conf->bufs, prev->bufs,
                                  (128 * 1024) / ngx_pagesize, ngx_pagesize);

    return NGX_CONF_OK;
}

static ngx_int_t
ngx_http_aes_filter_init(ngx_conf_t *cf)
{
    ngx_http_next_header_filter = ngx_http_top_header_filter;
    ngx_http_top_header_filter = ngx_http_aes_header_filter;

    ngx_http_next_body_filter = ngx_http_top_body_filter;
    ngx_http_top_body_filter = ngx_http_aes_body_filter;

    return NGX_OK;
}