/*
 * Public Domain 2007, Rene Herman
 *
 * gcc -W -Wall -DTEST -D_LARGEFILE64_SOURCE -o backup backup.c
 *
 */

#include <stdlib.h>

enum {
	DOS_EXTENDED   = 0x05,
	WIN98_EXTENDED = 0x0f,
	LINUX_EXTENDED = 0x85,
};

struct partition {
        unsigned char boot_ind;
	unsigned char __1[3];
        unsigned char sys_ind;
	unsigned char __2[3];
        unsigned int start;
        unsigned int size;
} __attribute__((packed));

struct entry {
	unsigned char flags;
	unsigned char type;
	unsigned short __1;
	unsigned long long start;
	unsigned int size;
} __attribute__((packed));

enum {
	ENTRY_FLAG_LOGICAL  = 0x01,
	ENTRY_FLAG_BOOTABLE = 0x80,
};

struct backup {
	unsigned char signature[8];
	unsigned short type;
	unsigned char heads;
	unsigned char sectors;
	unsigned char count;
	unsigned char __1[3];
	struct entry table[31];
} __attribute__((packed));

#define BACKUP_SIGNATURE "PARTBAK1"

enum {
	BACKUP_TYPE_MBR = 1,
};

struct backup backup = {
	.signature = BACKUP_SIGNATURE,
	.type	   = BACKUP_TYPE_MBR,
};

int is_extended(struct partition *entry)
{
	int ret = 0;

	switch (entry->sys_ind) {
	case DOS_EXTENDED:
	case WIN98_EXTENDED:
	case LINUX_EXTENDED:
		ret = 1;
	}
	return ret;
}

unsigned char *get_sector(unsigned int n);
void put_sector(unsigned char *sector);

int do_sector(unsigned int offset, unsigned int start)
{
	unsigned char *sector;
	struct partition *partition;
	struct partition *extended = NULL;

	start += offset;
	sector = get_sector(start);
	if (!sector) 
		return -1;

	if (sector[510] != 0x55 || sector[511] != 0xaa) {
		put_sector(sector);
		return -1;
	}

	partition = (struct partition *)(sector + 510 - 4 * sizeof *partition);
	do {
		struct entry *entry;

		if (!partition->size)
			continue;

		if (is_extended(partition)) {
			if (!offset)
				do_sector(partition->start, 0);
			else if (!extended)
				extended = partition;
			continue;
		}
		entry = backup.table + backup.count++;

		entry->flags = partition->boot_ind;
		if (offset)
			entry->flags |= ENTRY_FLAG_LOGICAL;
		entry->type  = partition->sys_ind;
		entry->start = start + partition->start;
		entry->size  = partition->size;

	} while (++partition < (struct partition *)(sector + 510));

	if (extended)
		do_sector(offset, extended->start);

	put_sector(sector);
	return 0;
}

#ifdef TEST
#include <stdio.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <fcntl.h>

#include <linux/hdreg.h>

int fd;

unsigned char *get_sector(unsigned int n)
{
	unsigned char *sector = malloc(512);

	if (!sector)
		return NULL;

	if (lseek64(fd, (off64_t)n << 9, SEEK_SET) < 0)
		goto free;

	if (read(fd, sector, 512) != 512) {
  free:		
		free(sector);
		sector = NULL;
	}
	return sector;
}

void put_sector(unsigned char *sector)
{
	free(sector);
}

int main(int argc, char *argv[])
{
	struct stat buf;
	struct hd_geometry geometry;

	if (argc != 2) {
		printf("%s <blkdev>\n", argv[0]);
		return EXIT_FAILURE;
	}
	fd = open(argv[1], O_RDONLY);
	if (fd < 0) {
		perror("open");
		return EXIT_FAILURE;
	}
	if (fstat(fd, &buf) < 0 || !S_ISBLK(buf.st_mode)) {
		perror("stat");
		return EXIT_FAILURE;
	}
	if (ioctl(fd, HDIO_GETGEO, &geometry) == 0) {
		backup.heads = geometry.heads;
		backup.sectors = geometry.sectors;
	}
	if (do_sector(0, 0) < 0)
		return EXIT_FAILURE;

	if (close(fd) < 0) {
		perror("close");
		return EXIT_FAILURE;
	}
	if (write(STDOUT_FILENO, &backup, sizeof backup) != sizeof backup) {
		perror("write");
		return EXIT_FAILURE;
	}
	return EXIT_SUCCESS;
}
#endif