#include <linux/bpf.h>
#define fuck_formater
#include <bpf/bpf_endian.h>
#include <bpf/bpf_helpers.h>

#define ARRAY_SIZE(a) sizeof(a) / sizeof(a[0])

#define ETH_HEADER_SIZE 14

#define MAC_SIZE 6

#define ETHER_TYPE_FIELD_OFFSET 12

#define ETHER_TYPE_IPV4 0x0800
#define ETHER_TYPE_IPV6 0x86DD

#define IPV4_TTL_OFFSET 8

#define IPV4_PROTOCOL_OFFSET 9
#define IPV4_CHECKSUM_OFFSET 10
#define IPV4_SRC_ADDR_OFFSET 12
#define IPV4_DST_ADDR_OFFSET 16

#define IPV6_NEXT_HEADER_OFFSET 6
#define IPV6_HOP_LIMIT_OFFSET 7
#define IPV6_SRC_ADDR_OFFSET 8
#define IPV6_DST_ADDR_OFFSET 24

#define IPV6_HEADER_SIZE 40

#define TCP_PROTO_ID 6

#define TCP_DST_PORT_OFFSET 2
#define TCP_SRC_PORT_OFFSET 0

#define TCP_SEQ_OFFSET 4

#define TCP_DATA_OFFSET_OFFSET 12
#define TCP_DATA_OFFSET_SHIFT 4

#define TCP_FLAGS_OFFSET 13

#define TCP_CHECKSUM_OFFSET 16

#define SYN_MASK 0b00000010
#define ACK_MASK 0b00010000

#define LOW_PORT 4000
const __u16 open_ports[] = {4444};
#define HIGH_PORT 5000

#define MAX_ITERATIONS 1024

char _license[] SEC("license") = "GPL";

__u64 time_seed = 0xffffffffffffffff;

#define DAY_LENGHT_IN_NS 86400000000000

#define DEFAULT_TTL 64

struct __attribute__((packed)) tcp_v4_pseudo_header { // BIG endian
  __u32 src_ip;
  __u32 dst_ip;
  __u8 zero_byte;
  __u8 protocol_number; // TCP_PROTO_ID
  __u16 tcp_segment_lenght;
};
struct __attribute__((packed)) tcp_v6_pseudo_header { // BIG endian
  __u128 src_ip;
  __u128 dst_ip;
  __u32 upper_layer_packet_lenght;
  __u8 zero[3];
  __u8 next_header_value; // TCP_PROTO_ID
};

__u32 mulberry32_at(__u32 seed, __u32 index) {
  __u32 z = seed + index * 0x6D2B79F5u;
  z = (z ^ (z >> 15)) * (z | 1u);
  z ^= z + (z ^ (z >> 7)) * (z | 61u);
  return z ^ (z >> 14);
}

static __always_inline const bool
binary_search(const __u16 *arr, const __u64 arr_len, const __u16 searched_for) {
  if (arr_len == 0) {
    return false;
  }
  if (searched_for < arr[0]) {
    return false;
  } else if (arr[0] == searched_for) {
    return true;
  } else {
    if (searched_for > arr[arr_len - 1]) {
      return false;
    }
    __u64 lover_bound_index = 0;
    __u64 upper_bound_index = arr_len - 1;
    for (__u64 i = 0; i < MAX_ITERATIONS; i++) {
      __u64 half_index =
          lover_bound_index + (upper_bound_index - lover_bound_index) / 2;
      if (searched_for < arr[half_index]) {
        upper_bound_index = half_index;
      } else if (searched_for > arr[half_index]) {
        lover_bound_index = half_index + 1;
      } else if (searched_for == arr[half_index]) {
        return true;
      }
    }
  }
  return false;
}

static __always_inline const __u8 *get_tcp_data(const __u16 EtherType,
                                                const __u8 *ip_packet_data,
                                                const __u8 *data_end) {
  if (EtherType == ETHER_TYPE_IPV4) {
    if (ip_packet_data + 1 > data_end) {
      return NULL;
    }

    // check protocol
    const __u8 *protocol_id_ptr = ip_packet_data + IPV4_PROTOCOL_OFFSET;
    if (protocol_id_ptr + 1 > data_end) {
      return NULL;
    }
    if (*protocol_id_ptr == TCP_PROTO_ID) {
      if (ip_packet_data + ((*ip_packet_data & 0b00001111) * 4) > data_end) {
        return NULL;
      }
      return ip_packet_data + ((*ip_packet_data & 0b00001111) * 4);
    }

  } else if (EtherType == ETHER_TYPE_IPV6) {
    // check  Next Header
    const __u8 *protocol_id_ptr = ip_packet_data + IPV6_NEXT_HEADER_OFFSET;
    if (protocol_id_ptr + 1 > data_end) {
      return NULL;
    }
    if (*protocol_id_ptr == TCP_PROTO_ID) {
      if (ip_packet_data + IPV6_HEADER_SIZE > data_end) {
        return NULL;
      }
      return ip_packet_data + IPV6_HEADER_SIZE;
    }
  }
  return NULL;
}

static __always_inline __u16 tcp_checksum_fold(__u32 sum) {
  sum = (sum & 0xFFFF) + (sum >> 16);
  sum = (sum & 0xFFFF) + (sum >> 16);
  sum = (sum & 0xFFFF) + (sum >> 16);
  return (__u16)~sum;
}

static __always_inline __u32 tcp_checksum_add(const void *data,
                                                const void *data_end,
                                                __u32 sum) {
  const __u8 *p = data;
  for (int i = 0; i < 740; i++) { // MTU/2, upper bound for the verifier
    if (p + 2 > (const __u8 *)data_end) {
      break;
    }
    sum += (p[0] << 8) | p[1];
    p += 2;
  }
  if (p + 1 <= (const __u8 *)data_end) {
    sum += p[0] << 8;
  }
  return sum;
}

SEC("faker")
int xdp_drop_prog(struct xdp_md *ctx) {
  if (bpf_ktime_get_ns() - time_seed > DAY_LENGHT_IN_NS) {
    time_seed = bpf_ktime_get_ns();
  }

  const __u8 *data_end = (__u8 *)(long)ctx->data_end;
  const __u8 *data = (__u8 *)(long)ctx->data;

  if (data + ETHER_TYPE_FIELD_OFFSET + sizeof(__u16) > data_end) {
    return XDP_PASS;
  }
  const __u16 EtherType = bpf_ntohs(*(__u16 *)(data + ETHER_TYPE_FIELD_OFFSET));

  const __u8 *ip_packet_data = data + ETH_HEADER_SIZE;
  if (ip_packet_data + 20 > data_end) { // it wont be smaller than this
    return XDP_PASS;
  }

  const __u8 *tcp_packet_data =
      get_tcp_data(EtherType, ip_packet_data, data_end);
  if (!tcp_packet_data) {
    return XDP_PASS;
  }

  if (tcp_packet_data + TCP_CHECKSUM_OFFSET + 2 > data_end) {
    return XDP_PASS;
  }

  __u16 *dst_port_ptr = (__u16 *)(tcp_packet_data + TCP_DST_PORT_OFFSET);
  const __u16 dst_port = bpf_ntohs(*dst_port_ptr);

  if (dst_port < LOW_PORT) {
    return XDP_PASS;
  }
  if (dst_port > HIGH_PORT) {
    return XDP_PASS;
  }
  bpf_printk("packet to port: %d\n", dst_port);

  if (binary_search(open_ports, ARRAY_SIZE(open_ports), dst_port)) {
    return XDP_PASS;
  }
  __u8 *tcp_flags_ptr = (__u8 *)(tcp_packet_data + TCP_FLAGS_OFFSET);
  const __u8 tcp_flags = *tcp_flags_ptr;

  // check SYN
  if (tcp_flags & SYN_MASK) {
    bpf_printk("packet is SYN\n");
    if (/*mulberry32_at((__u32)time_seed, dst_port) % 3 == 0*/1) {
      bpf_printk("DROPING!\n");
      __u16 *src_port_ptr = (__u16 *)(tcp_packet_data + TCP_SRC_PORT_OFFSET);
      const __u16 src_port = bpf_ntohs(*src_port_ptr);

      *src_port_ptr = bpf_htons(dst_port);
      *dst_port_ptr = bpf_htons(src_port);

      *tcp_flags_ptr = tcp_flags | ACK_MASK;

      *(__u16 *)(tcp_packet_data + TCP_CHECKSUM_OFFSET) = 0; // must be zero for calculation

      const __u8 tcp_data_offset =
          ((*(__u8 *)(tcp_packet_data + TCP_DATA_OFFSET_OFFSET)) >>
           TCP_DATA_OFFSET_SHIFT) *
          sizeof(__u32);

      __u16 tcp_checksum; // big endian

      {
        __u32 tmp_sum = 0;
        if (tcp_packet_data + tcp_data_offset > data_end) {
          return XDP_PASS;
        }
        if (EtherType == ETHER_TYPE_IPV4) {
          __u32 *src_ip_ptr = (__u32 *)(ip_packet_data + IPV4_SRC_ADDR_OFFSET);
          __u32 *dst_ip_ptr = (__u32 *)(ip_packet_data + IPV4_DST_ADDR_OFFSET);

          { // swap src and dst ip
            __u32 src_ip_BE = *src_ip_ptr;
            *src_ip_ptr = *dst_ip_ptr;
            *dst_ip_ptr = src_ip_BE;
          }

          *(__u8 *)(ip_packet_data + IPV4_TTL_OFFSET) = DEFAULT_TTL;

          { // calculate ipv4 header checksum
            if ((*ip_packet_data & 0b00001111) != 5) {
              return XDP_PASS;
            }
            if (ip_packet_data + 20 > data_end) {
              return XDP_PASS;
            }
            __u16 *ipv4_checksum_ptr =
                (__u16 *)(ip_packet_data + IPV4_CHECKSUM_OFFSET);
            *ipv4_checksum_ptr = 0; // must be zero for calculation
            __u32 new_checksum = 0;
            for (int i = 0; i < 10; i++) {
              if (ip_packet_data + i * 2 + 2 > data_end) {
                break;
              }
              new_checksum += bpf_ntohs(*(__u16 *)(ip_packet_data + i * 2));
            }
            new_checksum = (new_checksum & 0xFFFF) + (new_checksum >> 16);
            new_checksum = (new_checksum & 0xFFFF) + (new_checksum >> 16);
            *ipv4_checksum_ptr = bpf_htons((__u16)new_checksum);
          }

          struct tcp_v4_pseudo_header pseudo_header = {
              *src_ip_ptr, *dst_ip_ptr, 0, TCP_PROTO_ID,
              bpf_htons(data_end - tcp_packet_data)};

          tmp_sum = tcp_checksum_add(
              &pseudo_header,
              (const __u8 *)&pseudo_header + sizeof(pseudo_header), tmp_sum);
        } else {
          if (ip_packet_data + IPV6_HEADER_SIZE > data_end) {
            return XDP_PASS;
          }
          __u128 *src_ip_ptr =
              (__u128 *)(ip_packet_data + IPV6_SRC_ADDR_OFFSET);
          __u128 *dst_ip_ptr =
              (__u128 *)(ip_packet_data + IPV6_DST_ADDR_OFFSET);

          { // swap src and dst ip
            __u128 src_ip_BE = *src_ip_ptr;
            *src_ip_ptr = *dst_ip_ptr;
            *dst_ip_ptr = src_ip_BE;
          }

          *(__u8 *)(ip_packet_data + IPV6_HOP_LIMIT_OFFSET) = DEFAULT_TTL;

          struct tcp_v6_pseudo_header pseudo_header = {
              *(__u128 *)(ip_packet_data + IPV6_SRC_ADDR_OFFSET),
              *(__u128 *)(ip_packet_data + IPV6_DST_ADDR_OFFSET),
              bpf_htonl(data_end - tcp_packet_data),
              {0, 0, 0},
              TCP_PROTO_ID};

          tmp_sum = tcp_checksum_add(
              &pseudo_header,
              (const __u8 *)&pseudo_header + sizeof(pseudo_header), tmp_sum);
        }

        tmp_sum = tcp_checksum_add(tcp_packet_data, data_end, tmp_sum);
        tcp_checksum = bpf_htons(tcp_checksum_fold(tmp_sum));
      }
      *(__u16 *)(tcp_packet_data + TCP_CHECKSUM_OFFSET) = tcp_checksum;
      { // swap SRC an dst MAC
        __u8 src_mac[MAC_SIZE];
        __builtin_memcpy(src_mac, data + MAC_SIZE, MAC_SIZE);
        __builtin_memcpy((__u8 *)data + MAC_SIZE, data, MAC_SIZE);
        __builtin_memcpy((__u8 *)data, src_mac, MAC_SIZE);
      }
      return XDP_TX;
    }
  } else {
    bpf_printk("packet is  not SYN\n");
  }

  return XDP_PASS;
}