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

extern VALUE cLibnet;

static VALUE net_s_decode_udp(VALUE self, VALUE bytes) 
{
    VALUE                  udp_obj;
    VALUE                  str;
    VALUE                  payload;
    struct libnet_udp_hdr *udp;
    u_int16_t              uh_sport, uh_dport, uh_ulen, uh_sum;

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

    str = StringValue(bytes);

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

    udp = (struct libnet_udp_hdr *)RSTRING(str)->ptr;

    uh_sport = ntohs(udp->uh_sport);
    uh_dport = ntohs(udp->uh_dport);
    uh_ulen  = ntohs(udp->uh_ulen);
    uh_sum   = ntohs(udp->uh_sum);

    rb_iv_set(udp_obj, "@src_port", UINT2NUM(uh_sport));
    rb_iv_set(udp_obj, "@dst_port", UINT2NUM(uh_dport));
    rb_iv_set(udp_obj, "@length",   UINT2NUM(uh_ulen));
    rb_iv_set(udp_obj, "@checksum", UINT2NUM(uh_sum));

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

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

    return udp_obj;
}

static VALUE net_build_udp(int argc, VALUE *argv, VALUE self)
{
    libnet_t      *l;
    VALUE          udp, v;
    u_int16_t      sp, dp, len, sum;
    char          *payload   = NULL;
    u_int32_t      payload_s = 0;
    VALUE          ptag_obj;
    libnet_ptag_t  ptag = 0;

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

    Data_Get_Struct(self, libnet_t, l);

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

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

    v   = rb_funcall(udp, rb_intern("src_port"), 0);
    sp  = NUM2UINT(v);
    v   = rb_funcall(udp, rb_intern("dst_port"), 0);
    dp  = NUM2UINT(v);
    v   = rb_funcall(udp, rb_intern("length"), 0);
    len = NUM2UINT(v);
    v   = rb_funcall(udp, rb_intern("checksum"), 0);
    sum = NUM2UINT(v);

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

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

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

    ptag = libnet_build_udp(sp, dp, len, sum, (u_int8_t *)payload,
        payload_s, l, ptag);

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

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

    return udp;
}

void define_udp_methods()
{
    VALUE cUDP = rb_const_get(cLibnet, rb_intern("UDP"));

    rb_define_singleton_method(cUDP, "decode", net_s_decode_udp, 1);

    rb_define_method(cLibnet, "build_udp", net_build_udp, -1);
}