/* Copyright status: this file is in the public domain. */

/*
 * disk-check - check a disk for weirdness
 *
 * This program takes a disk device and fills it with distinctive data
 *  (data such that a block's contents uniquely say which block it is).
 *  Also has a mode which checks, instead of writing.
 *
 * The data pattern assumes the disk device is no larger than 2^64
 *  blocks (of 512 bytes each).  Since off_t is unlikely to be wider
 *  than 64 bits, there is already a significantly more restrictive
 *  size limit in place in practice.
 *
 * Basic usage (see below for details):
 *
 *	disk-check -w disk-device
 *		Writes distinctive data to the disk-device.
 *
 *	disk-check -c disk-device
 *		Reads from disk-device and checks that the data found
 *		matches what -w would write.
 *
 * Options:
 *
 *	-w
 *	-c
 *		Specify write or check operation, as above.
 *
 *	-s NNNN
 *		Specifies the disk's size, in 512-byte sectors.  For
 *		real disks (devices which support DIOCGDINFO and have
 *		minor device numbers that follow the usual pattern),
 *		this is unnecessary.  When DIOCGDINFO fails (or if the
 *		disk-device isn't a device special file), a warning is
 *		printed and the first write()/read() failure is taken
 *		to indicate the end of the device (the size determined
 *		is printed in this case).  Specifying a nonzero -s
 *		overrides all automated size determination and restores
 *		the usual behaviour on write/read errors.  (Specifying
 *		-s 0 is equivalent to not specifying -s at all, and can
 *		be used to override an earlier -s value.)
 *
 *	-d disk-device
 *		Like specifying a disk device unadorned, except it's
 *		unambiguous when the device name begins, or might
 *		begin, with a dash.
 *
 *	-r NNNN
 *		Specifies a restart of a crashed run.  Operates
 *		normally, except it assumes that all sectors before
 *		NNNN have already been handled.
 *
 *	-m STRING
 *		Specifies a magic string.  This replaces the data that
 *		would otherwise be written (-w) or expected (-c) at the
 *		beginning of each sector, for as many bytes as the
 *		string is long.  (This is designed to be used to ensure
 *		that data written by one run is not mistaken for data
 *		written by another run.)  To be useful, the STRING
 *		should be enough less than the sector size (512 bytes)
 *		to allow room for the sector serial number.  If this is
 *		given on write, it must be given, with the same string,
 *		when checking, or the compare will fail.
 *
 *	-e STRING
 *		On-disk data blocks are encrypted.  This is designed to
 *		defeat compression schemes that normally could be very
 *		effective on data as repetitive as the data generated
 *		by this program.  When this option is given, each
 *		block's data is generated as usual, but it is then
 *		encrypted.  Since the point is to avoid repetitive
 *		patterns, each block is encrypted with a different key.
 *		IDEA is used as a cryptosystem; the key for a block is
 *		based on STRING.  First, a 16-byte block is generated
 *		by feeding STRING through MD5.  The resulting 16-byte
 *		block then has the block number added to it to give the
 *		key for that block.  Blocks are, conceptually,
 *		encrypted immediately before being written and
 *		decrypted immediately after being read.  Each block
 *		still knows its block number, however; the block is
 *		computed as described, but then one or more bytes at
 *		the end are overwritten with the block number.  It is
 *		stored in raw base-256, using as few digits as possible
 *		given the size of the number.  These bytes are ordered
 *		most-significant-first, then have a byte suffixed to
 *		them whose value is the number of bytes needed for the
 *		number (not counting the byte-count suffix byte).  The
 *		result overwrites the end of the block (positioned so
 *		the byte-count byte replaces the last byte).
 *
 *		Of course, if -e is given when writing, it must also be
 *		given when reading, or the read test will fail.
 *
 *		This procedure is not intended to be good cryptographic
 *		practice in any sense.  It is simply intended to defeat
 *		compression attempts.
 *
 *	-D
 *	-E NNN
 *		As an aid to dealing with the data created by -e, -D
 *		and -E can be used to decrypt and encrypt,
 *		respectively, arbitrary data using the encryption
 *		algorithm outlined under -e.  (Again, this is not good
 *		cryptography; if you want real encryption, use a
 *		program designed for it.)  -E encrypts, as outlined
 *		under -e, given the block number NNN; -D decrypts.
 *		Since blocks know their block numbers, as outlined
 *		under -e, -D does not take a block number.  Each mode
 *		reads 512 bytes from stdin; -E writes 512 bytes to
 *		stdout, while -D writes less - usually 504 bytes, or
 *		496 for blocks above the 32 exabyte mark.  This is
 *		because the presence of the block number means the last
 *		IDEA block (or more than that if the block number, with
 *		its suffix, occupies more than 8 bytes, which happens
 *		at NNN = 2^56, which for 512-byte blocks means 2^65
 *		bytes, or 32EB) will decrypt to garbage.  This garbage
 *		is not printed.
 *
 *		The STRING comes from the -e option, which must be
 *		given.
 *
 *		If either of these options is given, -w, -c, and any
 *		disk-device argument are ignored.
 *
 * Except when -E or -D is given, a disk-device must be given, either
 *  with -d or without, and one of -w or -c must be given.
 *
 * If the disk-device is the single character `-' (regardless of how
 *  it's specified, ie, with -d or without), then disk-check will write
 *  to stdout for -w or read from stdin for -c.  -s must be used in
 *  this case; even if stdin/stdout is connected to a device that would
 *  normally auto-size, using this syntax suppresses that.
 */

#define NBLK 8192 /* sectors per chunk handled */

#include <md5.h>
#include <idea.h>
#include <stdio.h>
#include <fcntl.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <strings.h>
#include <sys/dkio.h>
#include <sys/ioctl.h>
#include <sys/disklabel.h>

extern const char *__progname;

typedef unsigned long long int ULLI;

#define OP_UNSET   1
#define OP_WRITE   2
#define OP_CHECK   3
#define OP_ENCRYPT 4
#define OP_DECRYPT 5
static int opmode = OP_UNSET;
static ULLI size = 0;
static int autosize = 0;
static const char *diskdev = 0;
static ULLI firstsect = 0;
static char *magic = 0;
static int pipemode = 0;
static int enc = 0;

static unsigned char keybase[16];
static int magiclen;
static int dfd;
static ULLI curblk;
#define DATASIZE (NBLK*512)
#ifdef DATA_ALIGN
static unsigned char data_[DATASIZE+DATA_ALIGN];
static unsigned char *data;
#else
static unsigned char data[DATASIZE];
#endif
static unsigned char rbuf[DATASIZE];
static volatile int gotsiginfo;
static FILE *siginfof;

static void set_key_base(const char *arg)
{
 void *h;

 h = md5_init();
 md5_process_bytes(h,arg,strlen(arg));
 md5_result(h,&keybase[0]);
}

static void handleargs(int ac, char **av)
{
 int skip;
 int errs;

 skip = 0;
 errs = 0;
 for (ac--,av++;ac;ac--,av++)
  { if (skip > 0)
     { skip --;
       continue;
     }
    if (**av != '-')
     { if (! diskdev)
	{ diskdev = *av;
	}
       else
	{ fprintf(stderr,"%s: extra argument `%s'\n",__progname,*av);
	  errs ++;
	}
       continue;
     }
    if (0)
     {
needarg:;
       fprintf(stderr,"%s: %s needs a following argument\n",__progname,*av);
       errs ++;
       continue;
     }
#define WANTARG() do { if (++skip >= ac) goto needarg; } while (0)
    if (!strcmp(*av,"-c"))
     { opmode = OP_CHECK;
       continue;
     }
    if (!strcmp(*av,"-w"))
     { opmode = OP_WRITE;
       continue;
     }
    if (!strcmp(*av,"-d"))
     { WANTARG();
       diskdev = av[skip];
       continue;
     }
    if (!strcmp(*av,"-e"))
     { WANTARG();
       enc = 1;
       set_key_base(av[skip]);
       continue;
     }
    if (!strcmp(*av,"-D"))
     { opmode = OP_DECRYPT;
       continue;
     }
    if (!strcmp(*av,"-E"))
     { WANTARG();
       opmode = OP_ENCRYPT;
       curblk = strtouq(av[skip],0,0);
       continue;
     }
    if (!strcmp(*av,"-s"))
     { WANTARG();
       size = strtouq(av[skip],0,0);
       continue;
     }
    if (!strcmp(*av,"-r"))
     { WANTARG();
       firstsect = strtouq(av[skip],0,0);
       continue;
     }
    if (!strcmp(*av,"-m"))
     { WANTARG();
       magic = av[skip];
       magiclen = strlen(magic);
       continue;
     }
#undef WANTARG
    fprintf(stderr,"%s: unrecognized option `%s'\n",__progname,*av);
    errs ++;
  }
 if (errs) exit(1);
 switch (opmode)
  { default:
       abort();
       break;
    case OP_UNSET:
       fprintf(stderr,"%s: must specify -w, -c, -D, or -E\n",__progname);
       exit(1);
       break;
    case OP_WRITE:
    case OP_CHECK:
       if (! diskdev)
	{ fprintf(stderr,"%s: must specify disk to operate on\n",__progname);
	  exit(1);
	}
       break;
    case OP_DECRYPT:
    case OP_ENCRYPT:
       if (! enc)
	{ fprintf(stderr,"%s: -D and -E require -e\n",__progname);
	  exit(1);
	}
       break;
  }
}

static void setup(void)
{
#ifdef DATA_ALIGN
 data = data_;
 if ((unsigned long int)data % DATA_ALIGN)
  { data += DATA_ALIGN - ((unsigned long int)data % DATA_ALIGN);
  }
#endif
 pipemode = !strcmp(diskdev,"-");
 if (pipemode && (size == 0))
  { fprintf(stderr,"%s: must use -s when giving `-' as device\n",__progname);
    exit(1);
  }
 switch (opmode)
  { case OP_WRITE:
       dfd = pipemode ? 1 : open(diskdev,O_WRONLY,0);
       break;
    case OP_CHECK:
       dfd = pipemode ? 0 : open(diskdev,O_RDONLY,0);
       break;
    default:
       abort();
       break;
  }
 if (dfd < 0)
  { fprintf(stderr,"%s: %s: %s\n",__progname,diskdev,strerror(errno));
    exit(1);
  }
 autosize = 0;
 if (size == 0)
  { struct disklabel lab;
    struct stat stb;
    fstat(dfd,&stb);
    switch (stb.st_mode & S_IFMT)
     { case S_IFCHR:
       case S_IFBLK:
	  if (ioctl(dfd,DIOCGDINFO,&lab) < 0)
	   { fprintf(stderr,"%s: note: %s: DIOCGDINFO: %s\n",__progname,diskdev,strerror(errno));
	     autosize = 1;
	   }
	  else
	   { if (lab.d_secsize == 512)
	      { size = lab.d_partitions[DISKPART(minor(stb.st_rdev))].p_size;
	      }
	     else
	      { fprintf(stderr,"%s: %s: labeled as using %d-byte sectors\n",__progname,diskdev,(int)lab.d_secsize);
		fprintf(stderr,"%s: attempting to compensate; recommend -s\n",__progname);
		size = lab.d_partitions[DISKPART(minor(stb.st_rdev))].p_size;
		if (lab.d_secsize & (lab.d_secsize-1))
		 { size = (size * lab.d_secsize) / 512;
		 }
		else if (lab.d_secsize < 512)
		 { size /= 512 / lab.d_secsize;
		 }
		else
		 { size *= lab.d_secsize / 512;
		 }
	      }
	   }
	  break;
       default:
	  fprintf(stderr,"%s: note: %s: not a device special file\n",__progname,diskdev);
	  autosize = 1;
	  break;
     }
  }
}

static void gen_key(unsigned char *kb, ULLI n)
{
 int i;

 bcopy(&keybase[0],kb,16);
 for (i=0;(i<16)&&n;i++)
  { n += kb[i];
    kb[i] = n & 0xff;
    n >>= 8;
  }
}

static void encrypt_block(unsigned char *b, ULLI n)
{
 unsigned char kb[16];
 unsigned char iv[8];
 IDEA_KEY k;
 int i;
 int j;

 gen_key(&kb[0],n);
 idea_setkey_e(&kb[0],&k);
 bzero(&iv[0],8);
 for (j=0;j<512;j+=8)
  { for (i=0;i<8;i++) iv[i] ^= b[j+i];
    idea_crypt(&iv[0],&k,&iv[0]);
    bcopy(&iv[0],b+j,8);
  }
 i = 510;
 while (n > 0)
  { b[i--] = n & 0xff;
    n >>= 8;
  }
 b[511] = 510 - i;
}

static int decrypt_block(unsigned char *b)
{
 unsigned char kb[16];
 unsigned char iv[8];
 unsigned char ct[8];
 IDEA_KEY k;
 int i;
 int j;
 ULLI n;
 int skipblks;

 skipblks = (b[511] + 8) >> 3;
 n = 0;
 for (i=511-b[511];i<511;i++) n = (n << 8) | b[i];
 gen_key(&kb[0],n);
 idea_setkey_d(&kb[0],&k);
 bzero(&iv[0],8);
 for (j=0;j<512;j+=8)
  { idea_crypt(b+j,&k,&ct[0]);
    for (i=0;i<8;i++) ct[i] ^= iv[i];
    bcopy(b+j,&iv[0],8);
    bcopy(&ct[0],b+j,8);
  }
 return(512-(skipblks*8));
}

static void fillblock(unsigned char *d, ULLI n)
{
 int i;

 sprintf(d,"%23llu",n);
 for (i=0;d[i]==' ';i++) d[i] = '.';
 i = 23;
 while (i < 256)
  { bcopy(d,d+i,i);
    i <<= 1;
  }
 bcopy(d,d+i,512-i);
 d[0] = '<';
 d[511] = '>';
 if (magic) bcopy(magic,d,magiclen);
 if (enc) encrypt_block(d,n);
}

static void filldata(void)
{
 int i;

 for (i=0;i<NBLK;i++) fillblock(&data[i*512],curblk+i);
}

static void print_humansize(FILE *to, ULLI nb)
{
 int p;

 p = 0;
 if (nb > ((~0ULL)>>10))
  { nb >>= 10;
    p = 1;
  }
 nb *= 512;
 while (nb >= 100000)
  { nb >>= 10;
    p ++;
  }
 fprintf(to,"%u",(unsigned int)nb);
 switch (p)
  { case 0: break;
    case 1: fprintf(to,"K"); break;
    case 2: fprintf(to,"M"); break;
    case 3: fprintf(to,"G"); break;
    case 4: fprintf(to,"T"); break;
    case 5: fprintf(to,"P"); break;
    case 6: fprintf(to,"E"); break;
    default: fprintf(to,"?%d",p); break;
  }
}

static void auto_sized(ULLI nblks, int offset)
{
 size = nblks;
 if (offset >= 512)
  { size += offset >> 9;
    offset &= 511;
  }
 printf("size: %llu sectors (",size);
 print_humansize(stdout,size);
 printf(")");
 if (offset) printf(" [+ %d leftover]",offset);
 printf("\n");
 autosize = 0;
}

static void siginfo(int sig __attribute__((__unused__)))
{
 gotsiginfo = 1;
}

static int siginfo_w(void *cookie __attribute__((__unused__)), const char *data, int len)
{
 static int fd = -1;

 switch (fd)
  { case -1:
       fd = open("/dev/tty",O_WRONLY,0);
       if (fd < 0)
	{ fd = -2;
	  fprintf(stderr,"%s: handling SIGINFO: open /dev/tty: %s\n",
		       __progname,strerror(errno));
    case -2:
	  break;
	}
    default:
       write(fd,data,len);
       break;
  }
 return(len);
}

static void setup_siginfo(void)
{
 struct sigaction sa;

 gotsiginfo = 0;
 sa.sa_handler = &siginfo;
 sigemptyset(&sa.sa_mask);
 sa.sa_flags = 0;
 sigaction(SIGINFO,&sa,0);
 siginfof = fwopen(0,&siginfo_w);
}

static void writeit(void)
{
 int n;
 int o;
 int t;
 int w;

 setup_siginfo();
 curblk = firstsect;
 if (! pipemode) lseek(dfd,curblk*512,SEEK_SET);
 while (autosize || (curblk < size))
  { filldata();
    w = NBLK;
    if (!autosize && (curblk+w > size)) w = size - curblk;
    w <<= 9;
    o = 0;
    while (o < w)
     { if (gotsiginfo)
	{ gotsiginfo = 0;
	  fprintf(siginfof,"sector %llu (",curblk+(o>>9));
	  print_humansize(siginfof,curblk+(o>>9));
	  fprintf(siginfof,")");
	  if (! autosize)
	   { fprintf(siginfof," of %llu (",size);
	     print_humansize(siginfof,size);
	     fprintf(siginfof,") = %7.3f%%",((curblk+(o>>9))*100.0)/size);
	   }
	  fprintf(siginfof,"\n");
	  fflush(siginfof);
	  continue;
	}
       t = w - o;
       n = write(dfd,&data[o],t);
       if (n == t) break;
       if (n < 0)
	{ if (errno == EINTR) continue;
	  if (autosize)
	   { auto_sized(curblk,o);
	     return;
	   }
	  n = errno;
	  fprintf(stderr,"%s: %s @%llu",__progname,diskdev,curblk);
	  if (o) fprintf(stderr,"%+d",o);
	  fprintf(stderr,": write: %s\n",strerror(n));
	  exit(1);
	}
       if (n == 0)
	{ fprintf(stderr,"%s: %s @%llu",__progname,diskdev,curblk);
	  if (o) fprintf(stderr,"%+d",o);
	  fprintf(stderr,": wrote nothing?\n");
	  if (autosize) auto_sized(curblk,o);
	  exit(1);
	}
       o += n;
     }
    curblk += w >> 9;
  }
}

static void checkit(void)
{
 int n;
 int o;
 int t;
 int r;
 int e;

 setup_siginfo();
 curblk = firstsect;
 if (! pipemode) lseek(dfd,curblk*512,SEEK_SET);
 e = 0;
 while (!e && (autosize || (curblk < size)))
  { filldata();
    r = NBLK;
    if (!autosize && (curblk+r > size)) r = size - curblk;
    r <<= 9;
    o = 0;
    while (o < r)
     { if (gotsiginfo)
	{ gotsiginfo = 0;
	  fprintf(siginfof,"sector %llu (",curblk+(o>>9));
	  print_humansize(siginfof,curblk+(o>>9));
	  fprintf(siginfof,")");
	  if (! autosize)
	   { fprintf(siginfof," of %llu (",size);
	     print_humansize(siginfof,size);
	     fprintf(siginfof,") = %7.3f%%",((curblk+(o>>9))*100.0)/size);
	   }
	  fprintf(siginfof,"\n");
	  fflush(siginfof);
	  continue;
	}
       t = r - o;
       n = read(dfd,&rbuf[o],t);
       if (n == t)
	{ o += t;
	  t = 0;
	  break;
	}
       if (n < 0)
	{ if (errno == EINTR) continue;
	  n = errno;
	  if (autosize) auto_sized(curblk,o);
	  fprintf(stderr,"%s: %s @%llu",__progname,diskdev,curblk);
	  if (o) fprintf(stderr,"%+d",o);
	  fprintf(stderr,": read: %s\n",strerror(n));
	  break;
	}
       if (n == 0)
	{ if (autosize) auto_sized(curblk,o);
	  break;
	}
       o += n;
     }
    if (bcmp(&rbuf[0],&data[0],o))
     { for (t=0;t<o;t++)
	{ if (rbuf[t] != data[t])
	   { printf("block %llu mismatch\n",curblk+(t>>9));
	     t = (t + 512) & ~511;
	   }
	}
     }
    if (o < r) return;
    curblk += r >> 9;
  }
}

static void just_decrypt(void)
{
 unsigned char blk[512];
 int n;

 n = fread(&blk[0],1,512,stdin);
 if (n != 512) exit(1);
 n = decrypt_block(&blk[0]);
 fwrite(&blk[0],1,n,stdout);
}

static void just_encrypt(void)
{
 unsigned char blk[512];
 int n;

 n = fread(&blk[0],1,512,stdin);
 if (n != 512) exit(1);
 encrypt_block(&blk[0],curblk);
 fwrite(&blk[0],1,512,stdout);
}

int main(int, char **);
int main(int ac, char **av)
{
 handleargs(ac,av);
 switch (opmode)
  { case OP_WRITE:
       setup();
       writeit();
       break;
    case OP_CHECK:
       setup();
       checkit();
       break;
    case OP_DECRYPT:
       just_decrypt();
       break;
    case OP_ENCRYPT:
       just_encrypt();
       break;
    default:
       abort();
       break;
  }
 exit(0);
}