/*
 * Copyright (c) 2023 Thomas Lindner <tom@dl6tom.de>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/param.h>

#include <sys/disklabel.h>
#include <sys/dkio.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <ufs/ffs/fs.h>
#include <ufs/ufs/dinode.h>
#include <ufs/ufs/dir.h>

#include <err.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <util.h>

int diskfd;
int filefd;
struct fs *sblock;

union dinode {
  struct ufs1_dinode dp1;
  struct ufs2_dinode dp2;
};
#define DIP(dp, field)                                                         \
  ((sblock->fs_magic == FS_UFS1_MAGIC) ? (dp)->dp1.field : (dp)->dp2.field)

typedef void (*process_fn)(union dinode *dp, void *datablock, size_t blksize,
                           off_t blkoff);

void usage(void);
void dump_inode(ino_t inum, union dinode *dp);
void dump_data(union dinode *dp, process_fn process_data,
               process_fn process_gap);
void dump_indirect(union dinode *dp, void *indirblock, int level,
                   off_t *filesize, process_fn process_data,
                   process_fn process_gap);

void process_directory(union dinode *dp, void *datablock, size_t blksize,
                       off_t blkoff);
void process_dirblk(int deleted, void *dirblk, size_t size);
void process_file(union dinode *dp, void *datablock, size_t blksize,
                  off_t blkoff);
void process_file_gap(union dinode *dp, void *datablock, size_t gapsize,
                      off_t blkoff);

void *xmalloc(size_t size);
void xpread(int fd, void *buf, size_t nbytes, off_t offset);
void xwrite(int fd, void *buf, size_t nbytes);
void xpwrite(int fd, void *buf, size_t nbytes, off_t offset);

int main(int argc, char **argv) {
#define MODE_ALL 1
#define MODE_INODE 2
#define MODE_DIRECTORY 4
#define MODE_FILE 8
  unsigned mode = 0;
  ino_t inode;

  int opt;
  const char *errstr;
  while ((opt = getopt(argc, argv, "adi:o:")) != -1) {
    switch (opt) {
    case 'a':
      if (mode & MODE_INODE) {
        usage();
      }
      mode |= MODE_ALL;
      break;
    case 'i':
      if (mode & MODE_ALL) {
        usage();
      }
      mode |= MODE_INODE;
      inode = strtonum(optarg, 2, UINT32_MAX, &errstr);
      if (errstr) {
        errx(1, "Invalid inode %s: %s", optarg, errstr);
      }
      break;
    case 'd':
      if (mode & MODE_FILE) {
        usage();
      }
      mode |= MODE_DIRECTORY;
      break;
    case 'o':
      if (mode & (MODE_ALL | MODE_DIRECTORY)) {
        usage();
      }
      mode |= MODE_FILE;
      if (!strcmp(optarg, "-")) {
        filefd = STDOUT_FILENO;
        break;
      }
      filefd = open(optarg, O_WRONLY | O_CREAT | O_EXCL,
                    S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);
      if (filefd == -1) {
        err(1, "Cannot open %s", optarg);
      }
      break;
    default:
      usage();
    }
  }
  if (!(mode & (MODE_ALL | MODE_INODE)) || optind + 1 != argc) {
    usage();
  }

  // open disk
  char *realpath;
  diskfd = opendev(argv[optind], O_RDONLY | O_NOFOLLOW, 0, &realpath);
  if (diskfd == -1) {
    err(1, "Cannot open %s (%s)", argv[optind], realpath);
  }

  // find superblock
  sblock = xmalloc(SBLOCKSIZE);
  off_t sblock_try[] = SBLOCKSEARCH;
  int i;
  for (i = 0; sblock_try[i] != -1; i++) {
    xpread(diskfd, sblock, SBLOCKSIZE, sblock_try[i]);
    if ((sblock->fs_magic == FS_UFS1_MAGIC ||
         (sblock->fs_magic == FS_UFS2_MAGIC &&
          sblock->fs_sblockloc == sblock_try[i])) &&
        sblock->fs_bsize <= MAXBSIZE && sblock->fs_bsize >= sizeof(struct fs)) {
      break;
    }
  }
  if (sblock_try[i] == -1) {
    errx(1, "Cannot find filesystem superblock");
  }

  void *inoblock = xmalloc(sblock->fs_bsize);
  if (mode & MODE_ALL) {
    // iterate over all inodes
    ino_t maxino = (ino_t)sblock->fs_ipg * sblock->fs_ncg;
    for (ino_t inum = 0; inum < maxino; inum++) {
      // read file system block
      if (!ino_to_fsbo(sblock, inum)) {
        daddr_t blkno = fsbtodb(sblock, ino_to_fsba(sblock, inum));
        xpread(diskfd, inoblock, sblock->fs_bsize, blkno * DEV_BSIZE);
      }
      // inode 0 is used as placeholder and 1 was used for bad blocks
      if (inum < 2) {
        continue;
      }
      union dinode *dp;
      if (sblock->fs_magic == FS_UFS1_MAGIC) {
        dp = (union dinode *)&(
            (struct ufs1_dinode *)inoblock)[ino_to_fsbo(sblock, inum)];
      } else {
        dp = (union dinode *)&(
            (struct ufs2_dinode *)inoblock)[ino_to_fsbo(sblock, inum)];
      }
      if (mode & MODE_DIRECTORY) {
        if ((DIP(dp, di_mode) & IFMT) == IFDIR) {
          dump_data(dp, process_directory, NULL);
        }
      } else {
        dump_inode(inum, dp);
      }
    }
  } else {
    daddr_t blkno = fsbtodb(sblock, ino_to_fsba(sblock, inode));
    xpread(diskfd, inoblock, sblock->fs_bsize, blkno * DEV_BSIZE);
    union dinode *dp;
    if (sblock->fs_magic == FS_UFS1_MAGIC) {
      dp = (union dinode *)&(
          (struct ufs1_dinode *)inoblock)[ino_to_fsbo(sblock, inode)];
    } else {
      dp = (union dinode *)&(
          (struct ufs2_dinode *)inoblock)[ino_to_fsbo(sblock, inode)];
    }
    if (mode & MODE_FILE) {
      dump_data(dp, process_file, process_file_gap);
      fprintf(stderr, "\r%" PRIu64 "/%" PRIu64 "\n", DIP(dp, di_size),
              DIP(dp, di_size));
    } else if (mode & MODE_DIRECTORY) {
      if ((DIP(dp, di_mode) & IFMT) != IFDIR) {
        errx(1, "Not a directory");
      }
      dump_data(dp, process_directory, NULL);
    } else {
      dump_inode(inode, dp);
    }
  }

  return 0;
}

void usage(void) {
  fprintf(stderr, "Usage: %s (-a|-i inode) [-d|-o outputfile] disk\n",
          getprogname());
  exit(1);
}

void dump_inode(ino_t inum, union dinode *dp) {
  char type;
  switch (DIP(dp, di_mode) & IFMT) {
  case IFIFO:
    type = 'p';
    break;
  case IFCHR:
    type = 'c';
    break;
  case IFDIR:
    type = 'd';
    break;
  case IFBLK:
    type = 'b';
    break;
  case IFREG:
    type = 'f';
    break;
  case IFLNK:
    type = 'l';
    break;
  case IFSOCK:
    type = 's';
    break;
  case IFWHT:
    type = 'w';
    break;
  default:
    type = 'u';
  }
  printf("inode:%" PRIu32 " type:%c mode:%" PRIo16 " nlink:%" PRId16
         " uid:%" PRIu32 " gid:%" PRIu32 " size:%" PRIu64 " atime:%" PRId64
         " mtime:%" PRId64 " ctime:%" PRId64 "\n",
         (uint32_t)inum, type, DIP(dp, di_mode) & ~IFMT, DIP(dp, di_nlink),
         DIP(dp, di_uid), DIP(dp, di_gid), DIP(dp, di_size), DIP(dp, di_atime),
         DIP(dp, di_mtime), DIP(dp, di_ctime));
}

void dump_data(union dinode *dp, process_fn process_data,
               process_fn process_gap) {
  char *datablock = xmalloc(sblock->fs_bsize);
  off_t filesize = DIP(dp, di_size);
  for (int i = 0; filesize > 0 && i < NDADDR;
       filesize -= sblock->fs_bsize, i++) {
    if (!DIP(dp, di_db[i])) {
      if (process_gap) {
        memset(datablock, 0, sblock->fs_bsize);
        process_gap(dp, datablock, sblock->fs_bsize,
                    DIP(dp, di_size) - filesize);
      }
      continue;
    }
    size_t blksize = sblksize(sblock, DIP(dp, di_size), i);
    daddr_t blkno = fsbtodb(sblock, DIP(dp, di_db[i]));
    xpread(diskfd, datablock, blksize, blkno * DEV_BSIZE);
    if (filesize < blksize) {
      blksize = filesize;
    }
    process_data(dp, datablock, blksize, DIP(dp, di_size) - filesize);
  }
  for (int i = 0; filesize > 0 && i < NIADDR; i++) {
    if (!DIP(dp, di_ib[i])) {
      size_t gapsize = NINDIR(sblock) * sblock->fs_bsize;
      for (int j = 0; j < i; j++) {
        gapsize *= NINDIR(sblock);
      }
      if (process_gap) {
        memset(datablock, 0, sblock->fs_bsize);
        process_gap(dp, datablock, gapsize, DIP(dp, di_size) - filesize);
      }
      filesize -= gapsize;
      continue;
    }
    daddr_t blkno = fsbtodb(sblock, DIP(dp, di_ib[i]));
    xpread(diskfd, datablock, sblock->fs_bsize, blkno * DEV_BSIZE);
    dump_indirect(dp, datablock, i, &filesize, process_data, process_gap);
  }
  free(datablock);
}

void dump_indirect(union dinode *dp, void *indirblock, int level,
                   off_t *filesize, process_fn process_data,
                   process_fn process_gap) {
  void *datablock = xmalloc(sblock->fs_bsize);
  for (int i = 0; *filesize > 0 && i < NINDIR(sblock); i++) {
    daddr_t fsblkno;
    if (sblock->fs_magic == FS_UFS1_MAGIC) {
      fsblkno = ((uint32_t *)indirblock)[i];
    } else {
      fsblkno = ((uint64_t *)indirblock)[i];
    }
    if (!fsblkno) {
      size_t gapsize = sblock->fs_bsize;
      for (int j = 0; j < level; j++) {
        gapsize *= NINDIR(sblock);
      }
      if (process_gap) {
        memset(datablock, 0, sblock->fs_bsize);
        process_gap(dp, datablock, gapsize, DIP(dp, di_size) - *filesize);
      }
      *filesize -= gapsize;
      continue;
    }
    size_t blksize = sblock->fs_bsize;
    daddr_t blkno = fsbtodb(sblock, fsblkno);
    xpread(diskfd, datablock, blksize, blkno * DEV_BSIZE);
    if (level) {
      dump_indirect(dp, datablock, level - 1, filesize, process_data,
                    process_gap);
    } else {
      if (*filesize < blksize) {
        blksize = *filesize;
      }
      process_data(dp, datablock, blksize, DIP(dp, di_size) - *filesize);
      *filesize -= sblock->fs_bsize;
    }
  }
  free(datablock);
}

void process_directory(union dinode *dp, void *datablock, size_t blksize,
                       off_t blkoff) {
  for (size_t j = 0; j < blksize; j += DIRBLKSIZ) {
    process_dirblk(0, datablock + j, MIN(blksize - j, DIRBLKSIZ));
  }
}

void process_dirblk(int deleted, void *dirblk, size_t size) {
  char *end = (char *)dirblk + size;
  for (struct direct *d = dirblk;
       (char *)&d->d_namlen < end && d->d_name + d->d_namlen - 1 < end;
       d = (struct direct *)((char *)d + d->d_reclen)) {
    char type;
    switch (d->d_type) {
    case DT_FIFO:
      type = 'p';
      break;
    case DT_CHR:
      type = 'c';
      break;
    case DT_DIR:
      type = 'd';
      break;
    case DT_BLK:
      type = 'b';
      break;
    case DT_REG:
      type = 'f';
      break;
    case DT_LNK:
      type = 'l';
      break;
    case DT_SOCK:
      type = 's';
      break;
    default:
      type = 'u';
    }
    char name[MAXNAMLEN + 1];
    memcpy(name, d->d_name, d->d_namlen);
    name[d->d_namlen] = '\0';
    printf("deleted:%d inode:%" PRIu32 " type:%c name:%s\n",
           deleted || !d->d_ino, d->d_ino, type, name);
    d->d_reclen = MIN(d->d_reclen, end - (char *)d);
    d->d_reclen &= ~3;
    if (!d->d_reclen) {
      if (!deleted) {
        warnx("Corrupted directory entry");
      }
      return;
    }
    if (d->d_reclen > DIRSIZ(NEWDIRFMT, d)) {
      process_dirblk(1, (char *)d + DIRSIZ(NEWDIRFMT, d),
                     d->d_reclen - DIRSIZ(NEWDIRFMT, d));
    }
  }
}

void process_file(union dinode *dp, void *datablock, size_t blksize,
                  off_t blkoff) {
  if (filefd == STDOUT_FILENO) {
    xwrite(filefd, datablock, blksize);
  } else {
    xpwrite(filefd, datablock, blksize, blkoff);
  }
  if (blkoff / (1024 * 1024) != (blkoff + blksize) / (1024 * 1024)) {
    fprintf(stderr, "\r%" PRIu64 "/%" PRIu64, blkoff + blksize,
            DIP(dp, di_size));
  }
}

void process_file_gap(union dinode *dp, void *datablock, size_t gapsize,
                      off_t blkoff) {
  if (filefd == STDOUT_FILENO) {
    for (int j = 0; j < gapsize / sblock->fs_bsize; j++) {
      xwrite(filefd, datablock, sblock->fs_bsize);
    }
  } else {
    if (ftruncate(filefd, blkoff + gapsize) == -1) {
      err(1, "ftruncate() failed");
    }
  }
  if (blkoff / (1024 * 1024) != (blkoff + gapsize) / (1024 * 1024)) {
    fprintf(stderr, "\r%" PRIu64 "/%" PRIu64, blkoff + gapsize,
            DIP(dp, di_size));
  }
}

void *xmalloc(size_t size) {
  void *p = malloc(size);
  if (!p) {
    err(1, "malloc() failed");
  }
  return p;
}

void xpread(int fd, void *buf, size_t nbytes, off_t offset) {
  ssize_t n = pread(fd, buf, nbytes, offset);
  if (n == -1) {
    err(1, "pread() failed");
  }
  if (n != nbytes) {
    errx(1, "Incomplete read");
  }
}

void xwrite(int fd, void *buf, size_t nbytes) {
  ssize_t n = write(fd, buf, nbytes);
  if (n == -1) {
    err(1, "write() failed");
  }
  if (n != nbytes) {
    errx(1, "Incomplete write");
  }
}

void xpwrite(int fd, void *buf, size_t nbytes, off_t offset) {
  ssize_t n = pwrite(fd, buf, nbytes, offset);
  if (n == -1) {
    err(1, "pwrite() failed");
  }
  if (n != nbytes) {
    errx(1, "Incomplete write");
  }
}