/*
 * Collect a set of DOS partition tables, the boot sector from each
 * partition, and output them in "densified compressed" format.
 */

#define _FILE_OFFSET_BITS 64

#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <linux/fs.h>

#define PARTGET_MAGIC	UINT64_C(0x3949831506e752b7)

/* A DOS partition table entry */
struct part_entry {
    uint8_t active_flag;        /* 0x80 if "active" */
    uint8_t start_head;
    uint8_t start_sect;
    uint8_t start_cyl;
    uint8_t ostype;
    uint8_t end_head;
    uint8_t end_sect;
    uint8_t end_cyl;
    uint32_t start_lba;
    uint32_t length;
} __attribute__ ((packed));

/* A DOS partition table */
struct part_table {
    char bootcode[440];
    uint32_t mbr_id;
    uint16_t pad;
    struct part_entry ptab[4];
    uint16_t sig;
} __attribute__((packed));

/* A data header */
struct data_header {
    uint64_t offset;
    uint64_t length;
};

#define SECTOR_SHIFT	9

static void dump_sector(FILE *dev, uint64_t lba)
{
    off_t offset = lba << SECTOR_SHIFT;
    char buf[1 << SECTOR_SHIFT];
    struct data_header hdr;

    if (fseeko(dev, offset, SEEK_SET))
	return;			/* Unable to seek */

    if (fread(&buf, 1, sizeof buf, dev) != sizeof buf)
	return;			/* Unable to read */

    hdr.offset = offset;
    hdr.length = sizeof buf;

    fwrite(&hdr, 1, sizeof hdr, stdout);
    fwrite(&buf, 1, sizeof buf, stdout);
}

static void get_ptab(FILE *dev, uint64_t lba, uint64_t base)
{
    off_t offset = lba << SECTOR_SHIFT;
    struct part_table ptab;
    struct data_header hdr;
    int i;

    if (fseeko(dev, offset, SEEK_SET))
	return;			/* Unable to seek */

    if (fread(&ptab, 1, sizeof ptab, dev) != sizeof ptab)
	return;			/* Unable to read */

    hdr.offset = offset;
    hdr.length = sizeof ptab;

    fwrite(&hdr, 1, sizeof hdr, stdout);
    fwrite(&ptab, 1, sizeof ptab, stdout);

    base = base ? base : lba;

    for (i = 0; i < 4; i++) {
	switch (ptab.ptab[i].ostype) {
	case 0x00:
	    /* Nothing */
	    break;
	case 0x05:
	case 0x0f:
	case 0x85:
	    /* Extended partition */
	    get_ptab(dev, base + ptab.ptab[i].start_lba, base);
	    break;
	default:
	    dump_sector(dev, lba + ptab.ptab[i].start_lba);
	    break;
	}
    }
}

int get_disk_size(FILE *dev)
{
    int fd = fileno(dev);
    struct data_header hdr;
    struct stat st;

    if (fstat(fd, &st))
	return -1;

    hdr.offset = PARTGET_MAGIC;

    if (S_ISREG(st.st_mode)) {
	hdr.length = st.st_size;
    } else if (S_ISBLK(st.st_mode)) {
	unsigned long blksize;

	if (!ioctl(fd, BLKGETSIZE64, &hdr.length)) {
	    /* done */
	} else if (!ioctl(fd, BLKGETSIZE, &blksize)) {
	    hdr.length = blksize << 9;
	} else {
	    return -1;
	}
    } else {
	errno = EINVAL;
	return -1;
    }
    
    fwrite(&hdr, 1, sizeof hdr, stdout);
    return 0;
}

int main(int argc, char *argv[])
{
    FILE *dev;

    if (argc != 2) {
	fprintf(stderr, "Usage: %s device > dump\n", argv[0]);
	exit(1);
    }

    dev = fopen(argv[1], "rb");
    if (!dev) {
	perror(argv[1]);
	exit(1);
    }

    if (get_disk_size(dev)) {
	perror(argv[1]);
	exit(1);
    }

    get_ptab(dev, 0, 0);
    return 0;
}