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

extern VALUE cLibnet;
static VALUE cARP;

static VALUE net_s_decode_arp(VALUE self, VALUE bytes) 
{
    VALUE                  arp_obj;
    VALUE                  str;
    VALUE                  payload;
    struct libnet_arp_hdr *arp;
    u_int16_t              ar_hrd;
    u_int16_t              ar_pro;
    u_int8_t               ar_hln;
    u_int8_t               ar_pln;
    u_int16_t              ar_op;
    int                    msg_len;
    int                    offset;

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

    str = StringValue(bytes);

    if (RSTRING(str)->len < LIBNET_ARP_H) {
        rb_raise(rb_eArgError, "string is too small to contain an ARP header");
    }

    arp = (struct libnet_arp_hdr *)RSTRING(str)->ptr;

    ar_hrd = ntohs(arp->ar_hrd);
    ar_pro = ntohs(arp->ar_pro);
    ar_hln = arp->ar_hln;
    ar_pln = arp->ar_pln;
    ar_op  = ntohs(arp->ar_op);

    rb_iv_set(arp_obj, "@htype",     UINT2NUM(ar_hrd));
    rb_iv_set(arp_obj, "@protocol",  UINT2NUM(ar_pro));
    rb_iv_set(arp_obj, "@hlen",      UINT2NUM(ar_hln));
    rb_iv_set(arp_obj, "@plen",      UINT2NUM(ar_pln));
    rb_iv_set(arp_obj, "@operation", UINT2NUM(ar_op));

    msg_len = LIBNET_ARP_H + (2 * ar_hln) + (2 * ar_pln);

    if (RSTRING(str)->len < msg_len) {
        rb_raise(rb_eArgError, "string is too small to contain an ARP message");
    }

    /* unpack the addresses */
    offset = LIBNET_ARP_H;
    rb_iv_set(arp_obj, "@sender_haddr", rb_str_new(RSTRING(str)->ptr + offset, ar_hln));
    offset += ar_hln;
    rb_iv_set(arp_obj, "@sender_paddr", rb_str_new(RSTRING(str)->ptr + offset, ar_pln));
    offset += ar_pln;
    rb_iv_set(arp_obj, "@target_haddr", rb_str_new(RSTRING(str)->ptr + offset, ar_hln));
    offset += ar_hln;
    rb_iv_set(arp_obj, "@target_paddr", rb_str_new(RSTRING(str)->ptr + offset, ar_pln));
    offset += ar_pln;

    /* unpack the payload */
    if (RSTRING(str)->len > msg_len) {
        payload = rb_str_new(RSTRING(str)->ptr + msg_len,
                RSTRING(str)->len - msg_len);
    }
    else {
        payload = Qnil;
    }

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

    return arp_obj;
}

static VALUE net_build_arp(int argc, VALUE *argv, VALUE self)
{
    libnet_t      *l;
    VALUE          arp_obj, v;
    u_int16_t      hrd, pro, op;
    u_int8_t       hln, pln;
    VALUE          sha, spa, tha, tpa; 
    char          *payload   = NULL;
    u_int32_t      payload_s = 0;
    VALUE          ptag_obj;
    libnet_ptag_t  ptag = 0;

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

    Data_Get_Struct(self, libnet_t, l);

    if (NIL_P(arp_obj)) {
        arp_obj = rb_class_new_instance(0, NULL, cARP);
        rb_yield(arp_obj);
    }

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

    v   = rb_funcall(arp_obj, rb_intern("htype"), 0);
    hrd = NUM2UINT(v);
    v   = rb_funcall(arp_obj, rb_intern("protocol"), 0);
    pro = NUM2UINT(v);
    v   = rb_funcall(arp_obj, rb_intern("hlen"), 0);
    hln = NUM2UINT(v);
    v   = rb_funcall(arp_obj, rb_intern("plen"), 0);
    pln = NUM2UINT(v);
    v   = rb_funcall(arp_obj, rb_intern("operation"), 0);
    op  = NUM2UINT(v);
    v   = rb_funcall(arp_obj, rb_intern("sender_haddr"), 0);
    sha = StringValue(v);
    v   = rb_funcall(arp_obj, rb_intern("sender_paddr"), 0);
    spa = StringValue(v);
    v   = rb_funcall(arp_obj, rb_intern("target_haddr"), 0);
    tha = StringValue(v);
    v   = rb_funcall(arp_obj, rb_intern("target_paddr"), 0);
    tpa = StringValue(v);

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

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

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

    ptag = libnet_build_arp(hrd, pro, hln, pln, op, RSTRING(sha)->ptr,
            RSTRING(spa)->ptr, RSTRING(tha)->ptr, RSTRING(tpa)->ptr,
            payload, payload_s, l, ptag);
            
    if (ptag == -1) {
        rb_raise(rb_eRuntimeError, libnet_geterror(l));
    }

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

    return arp_obj;
}

void define_arp_methods()
{
    cARP = rb_const_get(cLibnet, rb_intern("ARP"));

    rb_define_singleton_method(cARP, "decode", net_s_decode_arp, 1);

    rb_define_method(cLibnet, "build_arp", net_build_arp, -1);
}