#include <arpa/inet.h>
#include <ruby.h>
#include <libnet.h>

extern VALUE cLibnet;

static VALUE net_s_decode_vlan(VALUE self, VALUE bytes) 
{
    VALUE                       vlan_obj;
    VALUE                       str;
    VALUE                       payload;
    struct libnet_802_1q_hdr   *vlan;
    u_int16_t                   tpi;
    u_int16_t                   priority_c_vid;
    u_int16_t                   len;

    vlan_obj = rb_funcall(self, rb_intern("new"), 0);

    str = StringValue(bytes);

    if (RSTRING(str)->len < LIBNET_802_1Q_H) {
        rb_raise(rb_eArgError, "string is too small to contain a VLAN header");
    }

    vlan = (struct libnet_802_1q_hdr *)RSTRING(str)->ptr;

    tpi            = ntohs(vlan->vlan_tpi);
    priority_c_vid = ntohs(vlan->vlan_priority_c_vid);
    len            = ntohs(vlan->vlan_len);

    rb_iv_set(vlan_obj, "@dst", rb_str_new(vlan->vlan_dhost, ETHER_ADDR_LEN));
    rb_iv_set(vlan_obj, "@src", rb_str_new(vlan->vlan_shost, ETHER_ADDR_LEN));
    rb_iv_set(vlan_obj, "@tpi", UINT2NUM(tpi));
    rb_iv_set(vlan_obj, "@priority",
            UINT2NUM( (priority_c_vid >> 13) & LIBNET_802_1Q_PRIMASK) );
    rb_iv_set(vlan_obj, "@cfi",
            UINT2NUM( (priority_c_vid >> 12) & LIBNET_802_1Q_CFIMASK) );
    rb_iv_set(vlan_obj, "@id",
            UINT2NUM( priority_c_vid & LIBNET_802_1Q_VIDMASK) );
    rb_iv_set(vlan_obj, "@len", UINT2NUM(len));

    if (RSTRING(str)->len > LIBNET_802_1Q_H) {
        payload = rb_str_new(RSTRING(str)->ptr + LIBNET_802_1Q_H,
                RSTRING(str)->len - LIBNET_802_1Q_H);
    }
    else {
        payload = Qnil;
    }

    rb_iv_set(vlan_obj, "@payload", payload);

    return vlan_obj;
}

static VALUE net_build_vlan(int argc, VALUE *argv, VALUE self)
{
    libnet_t      *l;
    VALUE          vlan_obj, v;
    u_int8_t      *dst = NULL;
    u_int8_t      *src = NULL;
    u_int16_t      tpi;
    u_int8_t       priority;
    u_int8_t       cfi;
    u_int16_t      vlan_id;
    u_int16_t      len_proto;
    char          *payload = NULL;
    u_int32_t      payload_s = 0;
    VALUE          ptag_obj;
    libnet_ptag_t  ptag = 0;

    rb_scan_args(argc, argv, "01", &vlan_obj);

    Data_Get_Struct(self, libnet_t, l);

    if (NIL_P(vlan_obj)) {
        vlan_obj = rb_class_new_instance(0, NULL, rb_const_get(cLibnet, rb_intern("VLAN")));
        rb_yield(vlan_obj);
    }

    /* check that all required fields are set */
    rb_funcall(vlan_obj, rb_intern("check_packable"), 0);

    /* required parameters */
    v         = rb_funcall(vlan_obj, rb_intern("dst"), 0);
    dst       = (u_int8_t *)StringValuePtr(v);
    v         = rb_funcall(vlan_obj, rb_intern("src"), 0);
    src       = (u_int8_t *)StringValuePtr(v);
    v         = rb_funcall(vlan_obj, rb_intern("tpi"), 0);
    tpi       = NUM2UINT(v);
    v         = rb_funcall(vlan_obj, rb_intern("priority"), 0);
    priority  = NUM2UINT(v);
    v         = rb_funcall(vlan_obj, rb_intern("cfi"), 0);
    cfi       = NUM2UINT(v);
    v         = rb_funcall(vlan_obj, rb_intern("id"), 0);
    vlan_id   = NUM2UINT(v);
    v         = rb_funcall(vlan_obj, rb_intern("len"), 0);
    len_proto = NUM2UINT(v);

    /* optional parameters */
    if ((v = rb_funcall(vlan_obj, rb_intern("payload"), 0)) != Qnil) {
        v         = StringValue(v);
        payload   = RSTRING(v)->ptr;
        payload_s = RSTRING(v)->len;
    }

    ptag_obj = rb_iv_get(vlan_obj, "@ptag");

    if (!NIL_P(ptag_obj)) {
        ptag = NUM2LONG(ptag_obj);
    }

    ptag = libnet_build_802_1q(dst, src, tpi, priority, cfi, vlan_id, len_proto,
            (u_int8_t *)payload, payload_s, l, ptag);

    if (ptag == -1) {
        rb_raise(rb_eRuntimeError, libnet_geterror(l));
    }

    rb_iv_set(vlan_obj, "@ptag", LONG2NUM(ptag));

    return vlan_obj;
}

void define_vlan_methods()
{
    VALUE cVLAN = rb_const_get(cLibnet, rb_intern("VLAN"));

    rb_define_singleton_method(cVLAN, "decode", net_s_decode_vlan, 1);

    rb_define_method(cLibnet, "build_vlan", net_build_vlan, -1);
}