#include "socks5.h"

int socks5_handshake(int sock) {
    struct client_handshake ch = {
      .ver = 0x05,
      .nmethods = 0x01,
      .methods = {0x00}
    };

    int size = sizeof(uint8_t) * (2 + ch.nmethods);
    if (size != write(sock, &ch, size)) {
        fprintf(stderr, "partial/failed write\n");
        return -1;
    }

    struct server_handshake sh = {0};
    int err = read_entity(sock, &sh, sizeof(struct server_handshake));
    if (err == -1) {
        perror("sock5 handshake failed read");
        return -1;
    }
    if (ch.ver != sh.ver || sh.method != ch.methods[0]) {
      fprintf(stderr, "Protocol error: client asks for ver=%d, method=%d and server answers with ver=%d, method=%d\n",
              ch.ver, ch.methods[0], sh.ver, sh.method);
      return -1;
    }
    printf("[server_handshake] ver=%d, method=%d\n", sh.ver, sh.method);
    return 0;
}

int socks5_reply(int sock) {
    struct server_reply sr = {0};
    read_entity(sock, &sr, sizeof(uint8_t) * 4);
    switch(sr.atyp) {
        case ATYP_IPV4:
            if (read_entity(sock, sr.bind_addr.ipv4, sizeof(uint8_t) * 4) == -1)
                goto read_error;
            break;
        case ATYP_DOMAINNAME:
            if (read_entity(sock, &sr.bind_addr.dns.len, sizeof(uint8_t) * 4) == -1)
                goto read_error;
            if (read_entity(sock, sr.bind_addr.dns.str, sizeof(char) * sr.bind_addr.dns.len) == -1)
                goto read_error;
            break;
        case ATYP_IPV6:
            if (read_entity(sock, sr.bind_addr.ipv6, sizeof(uint8_t) * 16) == -1)
                goto read_error;
            break;
        default:
            fprintf(stderr, "Unsupported ATYP in server reply\n");
            return -128;
    }
    read_entity(sock, &sr.port, sizeof(uint16_t));

    if (sr.rep < 0 || sr.rep > 0x08) {
      fprintf(stderr, "Invalid reply field\n");
      return -128;
    }

    printf("[server_reply] ver=%d, rep=%s, atyp=%d, port=%d\n", sr.ver, rep_msg[sr.rep], sr.atyp, sr.port);
    return -sr.rep;

read_error:
    fprintf(stderr, "Unable to read ATYP\n");
    return -128;
}

int socks5_connect_dns(int sock, char* addr, uint16_t port) {
    char buffer[262];
    size_t domainLength = strlen(addr);
    if (domainLength > 255) {
        fprintf(stderr, "domain is too long\n");
        exit(EXIT_FAILURE);
    }
    struct client_request cr = { 0x05, CMD_CONNECT, 0x00, ATYP_DOMAINNAME, (uint8_t)domainLength, addr, htons(port) };
    size_t written = 0;
    fill_buffer(&written, buffer, &cr, 5*sizeof(uint8_t));
    fill_buffer(&written, buffer, cr.dst_addr, cr.dst_addr_len*sizeof(char));
    fill_buffer(&written, buffer, &cr.port, sizeof(uint16_t));

    if (written != write(sock, buffer, written)) {
        fprintf(stderr, "partial/failed write\n");
        return -1;
    }
}