/*  Copyright 2004 Paul Sladen <r31-acrdisk@paul.sladen.org>
**  GNU General Public License v2.  ABSOLUTELY NO WARRANTY.
**
**  2004-05-01  Figured out the MBR partition format;  Boot loader
**  occupies the first ~440 bytes and 4-rows of 16-byte partition
**  records start at byte 446 (0x1be).  LBA and Sectors are the
**  numbers of 512-byte sectors.
**
**  2004-05-03  Extended partitions (marked by ID==5) are a quasi-
**  linked-list.  Load the address of the extended partition from
**  the main MBR and save this for use as a base-reference.  Look
**  there and you find a partition table with two entries;  the first
**  is the partition you're after and the next (also marked ID==5)
**  is a link to the next partition in the linked list.  Offsets
**  are relative (additive) to the extended partition.
*/

#include <stdio.h>
#include <stdlib.h>
#include <error.h>
#include <assert.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <netinet/in.h>
#include "llseek.h"

static char usage[] = 
"Copyright 2003-2004 Paul Sladen <r31-acrdisk@paul.sladen.org>\n"
"GNU General Public License v2.  ABSOLUTELY NO WARRANTY.\n"
"\n"
"Usage: %s [options...]\n"
"\n"
"\t-d <file>\tUse <file> instead of '/dev/hda'\n"
"\t-n <file>\tUse <file> instead of '/dev/nvram'\n"
"\t-f <file>\tUse <file> instead of 'ibm1024r.bin'\n"
"\t-W\t\tEnable write operations (AT YOUR OWN RISK).\n"
"\n"
"1) You will need to create an empty partition the size of your RAM + 10MB,\n"
"with the partition type set to 'a0 ThinkPad Hibernation Partition'\n";

enum { 
  MBR_MAGIC = 0xaa55,
  MBR_LENGTH = 512,
  EXTENDED = 0x05,
  HIBERNATE = 0xa0,
  ZERO = 0,
  ACR_NVRAM_OFFSET = 70,
  ACR_NVRAM_LENGTH = 6,
  ACR_CREATED = 0xaa,
  SECTORS_TO_MB_SHIFT = 11,
  SECTORS_TO_KB_SHIFT = 1,
  EXTRA_BYTES = 10 * 1<<20,  /* 8MB + 768kB on my laptop, so call it 10MB */
};

typedef struct partition_entry {
  unsigned int padding:32;
  unsigned int id:8;
  unsigned int padding2:24;
  unsigned int lba:32;
  unsigned int sectors:32;
} partition_entry; 

typedef struct boot_record {
  unsigned char padding[446];
  partition_entry partition[4] __attribute__ ((packed));
  unsigned short magic;
} boot_record;

typedef struct acr_nvram {
  unsigned int lba:32;
  unsigned int flag:8;
  unsigned int checksum:8;
} acr_nvram;

enum {
  RAM = 0x42,
  LBA = 0x42,
  TIM = 0x42,
  CHK = 0x42
};

unsigned char const acr_default_disk_header[] = {
  'A',  'C',  'E',  'R',   ' ',  '0',  'V',  ' ',
  'S',  'U',  'S',  'P',   'E',  'N',  'D',  '.',
  '\0', RAM,  RAM,  RAM,   RAM,  0x00, 0x00, 0x00,
  0xaa, LBA,  LBA,  LBA,   LBA,  'A',  'C',  'R',
  '_',  '0',  'V',  ' ',   ' ',  'D',  'A',  'T',
  0xff, 0xff, 0x00, 0x00,  0xff, 0xff, 0x00, 0x00,
  0x00, 0x00, TIM,  TIM,   TIM,  TIM,  CHK
};

typedef struct acr_disk {
  char banner[17]; /* sixteen plus zero-terminated */
  unsigned long ram __attribute__ ((packed));
  char padding1[3];
  unsigned char flag;
  unsigned long lba __attribute__ ((packed));
  char filename[11];
  unsigned long _ffff[2] __attribute__ ((packed));
  char padding2[2];
  long timestamp __attribute__ ((packed));
  unsigned char checksum;
  char padding3[453];
  unsigned short version_major_be;
  unsigned short version_minor_be;
} acr_disk __attribute__ ((packed)); 

int
locate_partition(partition_entry *found, unsigned char requested_type, int fd,
		 unsigned long base, unsigned long offset)
{
  boot_record mbr;
  unsigned long extended = 0;
  int i;
  lloff_t aim = (lloff_t)(base + offset) * 512LL;

  /* Load and verify it /is/ a partition table */
  if((llseek(fd, aim, SEEK_SET) != aim) ||
     (read(fd, &mbr, MBR_LENGTH) != MBR_LENGTH && (errno = EFAULT)) ||
     (mbr.magic != MBR_MAGIC && (errno = EINVAL)))
    return errno;

  /* Search for the requested partition type, or note to follow extended partitions */
  for(i = 0; i < 4; i++)
    if( mbr.partition[i].id == EXTENDED )
      extended = mbr.partition[i].lba;
    else if( mbr.partition[i].id == requested_type )
      {
	mbr.partition[i].lba += base + offset;
	memcpy(found, &mbr.partition[i], sizeof(partition_entry));
	return errno = 0;
      }

  /* Got to the end and didn't find it! :-( */
  if( !extended )
    return errno = ENXIO;

  /* First item in extended partition [linked list], so save as base offset */
  if( extended && !base )
    {
      base = extended;
      extended = 0;
    }
  return locate_partition(found, requested_type, fd, base, extended);
}

int
acr_checksum (const char *s, int length)
{
  int checksum = 0;
  while(length-- > 0)
    checksum += *s++;
  return 0xff & -checksum;
}

int
read_nvram (acr_nvram *found, int fd)
{
  /* Try and read out the hibernation bytes from the CMOS */
  if((lseek(fd, ACR_NVRAM_OFFSET, SEEK_SET) != ACR_NVRAM_OFFSET) ||
     (read(fd, found, ACR_NVRAM_LENGTH) != ACR_NVRAM_LENGTH && (errno = EFAULT)))
    return errno;
  else if(found->flag == ACR_CREATED || !found->flag)
    return errno = acr_checksum((char *)found, ACR_NVRAM_LENGTH );
  else
    return errno = -1;
}

/* Call with  lba == zero  to clear the CMOS pointer */
int
write_nvram (int fd, unsigned long lba )
{
  acr_nvram hibernate = { 0, };

  /* Set the checksum; and the 0xaa flag if it really exists */
  if( hibernate.lba = lba )
    hibernate.flag = ACR_CREATED;
  hibernate.checksum = acr_checksum((char *)&hibernate, ACR_NVRAM_LENGTH - 1);
  
  /* Try and write out the hibernation bytes to the CMOS */
  if((lseek(fd, ACR_NVRAM_OFFSET, SEEK_SET) != ACR_NVRAM_OFFSET) ||
     ((write(fd, &hibernate, ACR_NVRAM_LENGTH) != ACR_NVRAM_LENGTH) && (errno = EINVAL)))
    return errno;
  else 
    return 0;
}

void
initialise_acr_disk (acr_disk *hibernate, unsigned long lba, unsigned long ram)
{
  memset( hibernate, 0x00, sizeof(acr_disk));
  memcpy( hibernate, acr_default_disk_header, sizeof(acr_default_disk_header));
  hibernate->ram = ram;
  hibernate->flag = ACR_CREATED;
  hibernate->lba = lba;
  hibernate->timestamp = time( 0 );
  hibernate->checksum = acr_checksum( &(hibernate->flag), ACR_NVRAM_LENGTH - 1 );
  hibernate->version_major_be = htons(1);
  hibernate->version_minor_be = htons(98);
}

int
read_acr_disk (acr_disk *hibernate, int fd, unsigned long sectors)
{
  lloff_t aim = (lloff_t)(sectors) * 512LL;

  /* Try and read out the hibernation bytes from the CMOS */
  if((llseek(fd, aim, SEEK_SET) != aim) ||
     (read(fd, hibernate, sizeof(acr_disk)) != sizeof(acr_disk) && (errno = EFAULT)))
    return errno;
  else if(memcmp(&hibernate->banner, acr_default_disk_header,
	    sizeof(hibernate->banner)))
    return errno = EINVAL;
  else
    return errno = (hibernate->checksum -
		    acr_checksum((char *)&hibernate->flag, ACR_NVRAM_LENGTH - 1 ));
}

int
write_acr_disk (int fd, const acr_disk *hibernate)
{
  lloff_t aim = (lloff_t)(hibernate->lba) * 512LL;

  printf("write_acr_disk(): aim == %lld, %lld, %lld\n", aim, aim >> 10, aim >> 20);

  /* NB: Should I validate and/or update the structure and checksum here? */

  /* Try and write out the hibernation bytes to the CMOS */
  if((llseek(fd, aim, SEEK_SET) != aim) ||
     (write(fd, hibernate, sizeof(acr_disk)) != sizeof(acr_disk)))
    {
      printf("current position (error) is: %lld\n", llseek(fd, 0, SEEK_CUR));
      return errno;
    }
  else 
    {
      printf("current position (okay) is: %lld\n", llseek(fd, 0, SEEK_CUR));
      return errno = 0;
    }
}

int
write_acr_firmware (int fd_disk, unsigned long lba, int fd_firmware, unsigned long *copied)
{
  lloff_t aim = (lloff_t)(lba) * 512LL;
  char buffer[512];
  ssize_t c = 1, r;
  int e;
  int remaining = *copied = 0;

  if((llseek(fd_disk, aim, SEEK_SET) != aim) ||
     (lseek(fd_firmware, 0, SEEK_SET)))
    return errno;

  while ( c )
    switch( c = read( fd_firmware, buffer, sizeof(buffer)) )
      {
      case -1: /* error */
	break;
      case 0: /* eof */
	memset( buffer, 0, sizeof(buffer));
	if ((r = write( fd_disk, buffer, remaining)))
	  {
	    printf("extra write said %d, ( remaining = %d)\n", r, remaining);
	    copied += r, errno = 0;
	  }
	break;
      default: /* 1->512 bytes */
	remaining = sizeof(buffer) - c;
	if( write( fd_disk, buffer, c) == c)
	  *copied += c;

    }
  return errno;
}

int
display_partition_entry(FILE *fp, const partition_entry *partition)
{
  return
    fprintf(fp,
	    "id\t\t: %#04x %d (%s)\n"
	    "lba\t\t: %#010x %d (%dMB)\n"
	    "sectors\t\t: %#010x %d (%dMB; %d blocks)\n",
	    partition->id, partition->id,
	    partition->id == HIBERNATE ? "ThinkPad Hibernation Partition" : "unknown",
	    partition->lba, partition->lba, partition->lba >> SECTORS_TO_MB_SHIFT,
	    partition->sectors, partition->sectors,
	    partition->sectors >> SECTORS_TO_MB_SHIFT,
	    partition->sectors >> SECTORS_TO_KB_SHIFT);
}

int
display_nvram(FILE *fp, const acr_nvram *hibernate)
{
  return
    fprintf(fp,
	    "lba\t\t: %#010x %d (%dMB)\n"
	    "flag\t\t: %#04x %d (%s)\n"
	    "checksum\t: %#04x %d (%s)\n",
	    hibernate->lba, hibernate->lba, hibernate->lba >> SECTORS_TO_MB_SHIFT,
	    hibernate->flag, hibernate->flag,
	    hibernate->flag == ACR_CREATED ? "CREATED" : "NOT CREATED",
	    hibernate->checksum, hibernate->checksum,
	    acr_checksum((char *)hibernate, ACR_NVRAM_LENGTH) ?
	    "INCORRECT" : "CORRECT");
}

int
display_acr_disk(FILE *fp, const acr_disk *hibernate)
{
  char banner[sizeof(hibernate->banner)+1];
  char filename[sizeof(hibernate->filename)+1];
  int major = ntohs(hibernate->version_major_be);
  int minor = ntohs(hibernate->version_minor_be);

  memcpy(banner, hibernate->banner, sizeof(hibernate->banner));
  memcpy(filename, hibernate->filename, sizeof(hibernate->filename));

  banner[sizeof(hibernate->banner)] = '\0';
  filename[sizeof(hibernate->filename)] = '\0';

  return
    fprintf(fp,
	    "banner\t\t: \"%.16s\"\n"
	    "ram\t\t: %#010x %d (%dMB)\n"
	    "created\t\t: %#04x %d (%s)\n"
	    "lba\t\t: %#010x %d (%dMB)\n"
	    "filename\t: \"%.11s\"\n"
	    "timestamp\t: %#010x %d (%.24s)\n"
	    "checksum\t: %#04x %d (%s)\n"
	    "version\t\t: %#04x %#04x (%d.%d)\n",
	    banner,
	    hibernate->ram, hibernate->ram, hibernate->ram >> 20,
	    hibernate->flag, hibernate->flag, 
	    hibernate->flag == ACR_CREATED ? "CREATED" : "NOT CREATED",
	    hibernate->lba, hibernate->lba, hibernate->lba >> SECTORS_TO_MB_SHIFT,
	    filename,
	    hibernate->timestamp, hibernate->timestamp, ctime(&hibernate->timestamp),
	    hibernate->checksum, hibernate->checksum,
	    (acr_checksum((char *)&hibernate->flag, ACR_NVRAM_LENGTH-1)
	     == hibernate->checksum) ? "CORRECT" : "INCORRECT", 
	    major, minor, major, minor);
}

int
main (int argc, char **argv)
{
  int fd_disk, fd_nvram, fd_firmware, r, c, mode = O_RDONLY;
  partition_entry partition = { 0, };
  acr_nvram cmos = { 0, };
  acr_disk hibernate = { 0, };
  char
    *disk_device = "/dev/hda",
    *nvram_device = "/dev/nvram",
    *firmware_device = "ibm1024r.bin";
  unsigned long size, total = 0;

  assert(sizeof(partition_entry) == 16);
  assert(sizeof(boot_record) == 512);

  while((c = getopt(argc, argv, "Wd:n:f:")) != -1 )
    switch(c)
      {
      case 'W':
	mode = O_RDWR;
	printf("Enabling WRITE operations\n");
	break;
      case 'd':
	disk_device = strdup(optarg);
	break;
      case 'n':
	nvram_device = strdup(optarg);
	break;
      case 'm':
	firmware_device = strdup(optarg);
	break;
      default:
	printf(usage, *argv);
	exit(EXIT_SUCCESS);
      }
  
  fprintf(stderr, "%s: see '%s -h' or the manpage for assistance\n", *argv, *argv);

  if( -1 == (fd_disk = open(disk_device, mode)))
    perror( disk_device ), exit(1);
  if( -1 == (fd_nvram = open(nvram_device, mode)))
    perror( nvram_device ), exit(1);
  if( mode != O_RDONLY)
    if( -1 == (fd_firmware = open(firmware_device, O_RDONLY)))
      perror( firmware_device ), exit(1);

#if 0
  if( r = write_nvram(fd_nvram, 0))
    perror("main(): write_nvram() backup said it failed because of"), exit(1);
  exit(0);
#endif
#if 0
  //  74 b1  5a 03 
  if( r = write_nvram(fd_nvram, 0x035ab174))
    perror("main(): write_nvram() backup said it failed because of"), exit(1);
  exit(0);
#endif

  printf("Trying to parse a 512-byte MBR chain from \"%s\"\n", disk_device);
  if( r = locate_partition(&partition, HIBERNATE, fd_disk, 0, 0))
    perror("main(): locate_partition() said it failed because of"),
      fprintf(stderr, "Have you created a partition with type == 0a ? (see: -h)\n"),
      exit(1);

  display_partition_entry(stdout, &partition);

  size = partition.sectors * 512;
  //  size -= EXTRA_BYTES;  /* Space for Video RAM and extras */
  // size &= ~0 << 20;  /* Round down to nearest MB */

  printf("\n** Initialising with RAM size = %d (%dMB), LBA offset = %d (~%dMB)\n",
	 size, size >> 20, partition.lba, partition.lba >> 11);

  if( size <= (16 * (1<<20)))
    fprintf(stderr,
	    "Are you sure you want to create a hibernation partition less than 16MB?\n"
	    "Comment this out in the source if you do...\n"),
      exit(1);

  initialise_acr_disk(&hibernate, partition.lba, size);

  printf("\nTrying to WRITE a 512-byte hibernation header to \"%s\"\n", disk_device);

  if(mode == O_RDONLY)
    fprintf(stderr, "Skipping.  Use  -W  to enable WRITE operations\n");
  else
    if( r = write_acr_disk(fd_disk, &hibernate))
      perror("main(): write_acr_disk() said it failed because of"), exit(1);

  printf("\nTrying to WRITE a 431kB hibernation firmware image to \"%s\"\n", disk_device);

  if(mode == O_RDONLY)
    fprintf(stderr, "Skipping.  Use  -W  to enable WRITE operations\n");
  else
    if( r = write_acr_firmware(fd_disk, partition.lba + 1, fd_firmware, &total))
      perror("main(): write_acr_firmware() said it failed because of"), exit(1);

  printf("firmware total = %d\n", total);

  printf("\nTrying to WRITE a 6-byte hibernation offset to \"%s\"\n", nvram_device);
  if(mode == O_RDONLY)
    fprintf(stderr, "Skipping.  Use  -W  to enable WRITE operations\n");
  else
    if( r = write_nvram(fd_nvram, partition.lba))
      perror("main(): write_nvram() said it failed because of"), exit(1);

  printf("\nTrying to parse a 6-byte hibernation offset from \"%s\"\n", nvram_device);
  if( r = read_nvram(&cmos, fd_nvram) )
    perror("main(): read_nvram() said it failed because of"), exit(1);

  display_nvram(stdout, &cmos);

  printf("\nTrying to parse a 512-byte hibernation header from \"%s\"\n", disk_device);
  if( r = read_acr_disk(&hibernate, fd_disk, cmos.lba))
    perror("main(): read_acr_disk() said it failed because of"), exit(1);

  display_acr_disk(stdout, &hibernate);

  close(fd_disk);
  close(fd_nvram);
  close(fd_firmware);
}