#include <ruby.h>
#include <libnet.h>

extern VALUE cLibnet;

static VALUE net_s_decode_tcp(VALUE self, VALUE bytes) 
{
    VALUE                  tcp_obj;
    VALUE                  str;
    VALUE                  payload;
    struct libnet_tcp_hdr *tcp;
    u_int16_t              th_sport, th_dport;
    u_int32_t              th_seq, th_ack;
    u_int8_t               th_off, th_x2, th_flags;
    u_int16_t              th_win, th_sum, th_urp;

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

    str = StringValue(bytes);

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

    tcp = (struct libnet_tcp_hdr *)RSTRING(str)->ptr;

    th_sport = ntohs(tcp->th_sport); 
    th_dport = ntohs(tcp->th_dport); 
    th_seq   = ntohl(tcp->th_seq);
    th_ack   = ntohl(tcp->th_ack);
    th_off   = tcp->th_off;
    th_x2    = tcp->th_x2 << 3;
    th_x2   |= (tcp->th_flags & 0xe0) >> 5;
    th_flags = tcp->th_flags & 0x1f;
    th_win   = ntohs(tcp->th_win);
    th_sum   = ntohs(tcp->th_sum);
    th_urp   = ntohs(tcp->th_urp);

    rb_iv_set(tcp_obj, "@src_port",    UINT2NUM(th_sport));
    rb_iv_set(tcp_obj, "@dst_port",    UINT2NUM(th_dport));
    rb_iv_set(tcp_obj, "@seq",         ULONG2NUM(th_seq));
    rb_iv_set(tcp_obj, "@ack",         ULONG2NUM(th_ack));
    rb_iv_set(tcp_obj, "@data_offset", UINT2NUM(th_off));
    rb_iv_set(tcp_obj, "@reserved",    UINT2NUM(th_x2));
    rb_iv_set(tcp_obj, "@control",     UINT2NUM(th_flags));
    rb_iv_set(tcp_obj, "@window",      UINT2NUM(th_win));
    rb_iv_set(tcp_obj, "@checksum",    UINT2NUM(th_sum));
    rb_iv_set(tcp_obj, "@urgent",      UINT2NUM(th_urp));

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

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

    return tcp_obj;
}

static VALUE net_build_tcp(int argc, VALUE *argv, VALUE self)
{
    libnet_t      *l;
    VALUE          tcp, v;
    u_int16_t      sp, dp;
    u_int32_t      seq, ack;
    u_int8_t       control;
    u_int16_t      win, sum, urg;
    char          *payload   = NULL;
    u_int32_t      payload_s = 0;
    VALUE          ptag_obj;
    libnet_ptag_t  ptag = 0;

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

    Data_Get_Struct(self, libnet_t, l);

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

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

    v       = rb_funcall(tcp, rb_intern("src_port"), 0);
    sp      = NUM2UINT(v);
    v       = rb_funcall(tcp, rb_intern("dst_port"), 0);
    dp      = NUM2UINT(v);
    v       = rb_funcall(tcp, rb_intern("seq"), 0);
    seq     = NUM2ULONG(v);
    v       = rb_funcall(tcp, rb_intern("ack"), 0);
    ack     = NUM2ULONG(v);
    v       = rb_funcall(tcp, rb_intern("control"), 0);
    control = NUM2UINT(v);
    v       = rb_funcall(tcp, rb_intern("window"), 0);
    win     = NUM2UINT(v);
    v       = rb_funcall(tcp, rb_intern("checksum"), 0);
    sum     = NUM2UINT(v);
    v       = rb_funcall(tcp, rb_intern("urgent"), 0);
    urg     = NUM2UINT(v);

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

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

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

    ptag = libnet_build_tcp(sp, dp, seq, ack, control, win, sum, urg,
        LIBNET_TCP_H + payload_s, (u_int8_t *)payload, payload_s, l, ptag);

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

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

    return tcp;
}

void define_tcp_methods()
{
    VALUE cTCP = rb_const_get(cLibnet, rb_intern("TCP"));

    rb_define_singleton_method(cTCP, "decode", net_s_decode_tcp, 1);

    rb_define_method(cLibnet, "build_tcp", net_build_tcp, -1);
}