// gcc -o level11 -no-pie -fstack-protector-all level11.c
#include <arpa/inet.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

const uint16_t port = 4011;

void win(int fd) {
  // Connect stdin and stdout to the client socket,
  // so they can interact with the shell.
  dup2(fd, STDIN_FILENO);
  dup2(fd, STDOUT_FILENO);
  dup2(fd, STDERR_FILENO);

  char *argv[2];
  argv[0] = "/bin/sh";
  argv[1] = NULL;
  execve(argv[0], argv, NULL);
}

void prompt(int fd) {
  char buffer[32];
  send(fd, "What do you want to talk about?\n", 32, 0);
  recv(fd, buffer, 120, 0);
}

void vuln(int fd) {
  prompt(fd);
  send(fd, "Bye.\n", 5, 0);
}

// forking socket server with help from
// https://github.com/pwning/docs/blob/master/fork_accept.c
int main(int argc, char **argv) {
  // Setting the SIGCHLD handler to SIG_IGN prevents child
  // processes from becoming zombies (so you do not need to
  // call wait() on them).
  if (signal(SIGCHLD, SIG_IGN) == SIG_ERR) {
    fputs("Failed to set SIGCHLD handler.", stderr);
    return 1;
  }

  // Create server socket.
  int server_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
  if (server_sock < 0) {
    perror("socket");
    return 1;
  }

  // Set SO_REUSEADDR. Otherwise, if the server crashes for
  // any reason, you will have to wait for sockets to time
  // out before you can reuse the port.
  int opt = 1;
  if (setsockopt(server_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) !=
      0) {
    perror("setsockopt");
    return 1;
  }

  struct sockaddr_in listen_addr = {0};
  bzero((char *)&listen_addr, sizeof(listen_addr));
  listen_addr.sin_family = AF_INET;
  listen_addr.sin_port = htons(port);
  if (inet_pton(AF_INET, "127.0.0.1", &listen_addr.sin_addr) <= 0) {
    perror("inet_pton");
    return 1;
  }

  if (bind(server_sock, (struct sockaddr *)&listen_addr, sizeof(listen_addr)) !=
      0) {
    perror("bind");
    return 1;
  }

  if (listen(server_sock, 5) != 0) {
    perror("listen");
    return 1;
  }

  int client_sock;
  pid_t child_pid;
  while (1) {
    client_sock = accept(server_sock, NULL, NULL);
    if (client_sock < 0) {
      perror("accept");
      continue;
    }

    child_pid = fork();
    if (!child_pid) {
      // Avoid tons of long-running processes sticking around.
      alarm(30);

      // If you do not close the socket fd, someone who
      // exploits the service could call accept() on it and
      // hijack other people's connections.
      close(server_sock);

      // Call the vulnerable code with the client socket.
      vuln(client_sock);
      close(client_sock);
      return 0;
    } else {
      // If you forget to close the client fd, you could run
      // out of file descriptors.
      close(client_sock);
    }
  }

  return 0;
}
