imager.cpp

sleuthit-2.09 一个磁盘的工具集

#include "config.h"
#include "afflib.h"
#include "afflib_i.h"
#include "imager.h"
#include "aimage.h"
#include "ident.h"
#include "gui.h"

#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <err.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <fcntl.h>
#include <string.h>
#include <signal.h>
#include <assert.h>
#include <errno.h>

#ifdef HAVE_TERM_H
#include <term.h>
#endif

#ifdef HAVE_NCURSES_TERM_H
#include <ncurses/term.h>
#endif



/*
 * imager.cpp:
 * The C++ imaging library.
 */

imager::imager()
{
    allow_regular = false;
    total_segments_written = 0;
    total_sectors_read=0;
    total_bytes_read = 0;
    total_bytes_written = 0;
    total_blank_sectors = 0;
    
    callback_bytes_to_write = 0;
    callback_bytes_written = 0;

    imaging = false;
    imaging_failed = false;

    logfile = 0;

    last_sector_read = 0;	// sector number
    bad_sectors_read = 0;
    af = 0;
    fout = 0;				// 
    hash_invalid = false;		// make true to avoid hash calculation

    memset(cmd_attach,0,sizeof(cmd_attach));
    memset(cmd_detach,0,sizeof(cmd_detach));

    scsi_bus = -1;
    scsi_tid = -1;
    scsi_lun = -1;
    scsi_pass = -1;
    ata_dev = -1;
    
    memset(device_model,0,sizeof(device_model));
    memset(serial_number,0,sizeof(serial_number));
    memset(firmware_revision,0,sizeof(firmware_revision));

    in     = -1;
    in_pos = 0;
    sector_size = 0;
    total_sectors = 0;
    maxreadblocks = 0;

    af = 0;
    fout = 0;

    memset(fname_raw,0,sizeof(fname_raw));
    memset(fname_aff,0,sizeof(fname_aff));
    memset(infile,0,sizeof(infile));

    memset(final_md5,0,sizeof(final_md5));
    memset(final_sha1,0,sizeof(final_sha1));
    hash_invalid = false;

    last_sector_read = 0;
    last_sectors_read = 0;

    seek_on_output = false;
    retry_count = 0;

    buf = 0;
    bufsize = 512;			// good guess
    memset(blank_sector,0,sizeof(blank_sector));
    partial_sector_left  = 0;
    partial_sector_blank = false;

    bad_sectors_read = 0;
    consecutive_read_errors = 0;
    consecutive_read_error_regions = 0;
    error_recovery_phase = 0;
    last_direction = 0;

    /* error recovery */
}


void imager::write_data(unsigned char *buf,uint64 offset,int len)
{
    /* if this is supposed to be bad data, make sure that it is properly bad... */
    if(opt_debug==99){
	printf("imager::write_data(buf/x=%p,offset=%qd len=%d buf=%s\n",buf,offset,len,buf);
	if(offset%sector_size != 0){
	    err(1,"huh? offset mod %d = %d\n",sector_size,(int)offset%sector_size);
	}
    }

    if(!hash_invalid){
	/* Update hash functions. */
	MD5_Update(&md5,buf,len);
	SHA1_Update(&sha,buf,len);
    }

    /* Count the number of blank sectors.
     */
    /* First, see if there is a partial blank sector that we are still processing... */
    int len_left = len;
    while(len_left>0 && partial_sector_left>0){
	if(buf[len-len_left] != 0){
	    partial_sector_blank = false; // it's no longer blank
	}
	len_left--;
	partial_sector_left--;
    }
    if(partial_sector_left==0){ // we reached the end of the partial sector
	if(partial_sector_blank==true){	// and the sector is blank!
	    total_blank_sectors++;
	}
    }
    /* Is it possible to look for full sectors? */
    while(len_left > sector_size){
	if(memcmp(buf+(len-len_left),blank_sector,sector_size)==0){
	    total_blank_sectors++;
	}
	len_left -= sector_size;
    }

    if(partial_sector_left==0 && len_left>0){ // some left, so we have a new partial sector
	partial_sector_left  = sector_size;
	partial_sector_blank = true;
    }

    /* If anything is left, do the partial sector */
    while(len_left>0){
	if(buf[len-len_left] != 0){
	    partial_sector_blank = false; // no longer blank
	}
	len_left--;
	partial_sector_left--;
    }


    /* Write it out and carry on... */
    if(af){
	if(offset) af_seek(af,offset,SEEK_SET);
	if(af_write(af,buf,len)!=len){
	    perror("af_write");	// this is bad
	    af_close(af);	// try to gracefully recover
	    fprintf(stderr,"\r\n");
	    fprintf(stderr,"Imaging terminated because af_write failed.\n\r");
	    exit(1);
	}
    }
    if(fout){
	if(offset) fseeko(fout,offset,SEEK_SET);
	if(fwrite(buf,1,len,fout)!=(unsigned)len){
	    perror("fwrite");	// this is also bad
	    fclose(fout);	// get what we can out to disk
	    fprintf(stderr,"Imaging terminated because fwrite() failed.\n");
	    fprintf(stderr,"Correct error condition and re-run aimage.\n");
	    fprintf(stderr,"You may be able to carry on from where you left off.\n");
	    exit(1);
	}
    }
    total_bytes_written   += len;
}

void imager::status()
{
    if(opt_quiet==0 && opt_silent==0){
	my_refresh(this,0);			// just refresh; most status is done by AFF callback
    }
}

/****************************************************************
 *** isleep(): An informative sleep...                        ***
 ****************************************************************/
void isleep(int s)
{
  printf("isleep %d\n",s);
  for(int i=0;i<s;i++){
    printf("\rSleeping for %d seconds; %d left...",s,s-i);
    fflush(stdout);
    sleep(1);
  }
  printf("\r%50s\r","");
  fflush(stdout);
}



/*
 * open_dev(char outdev[MAXPATHLEN],char *indev)
 * Try to open the friendly device name.
 * If successful, return the actual device in outdev and the FD.
 * 
 * If the device can be detected but not mounted (common with some
 * broken IDE drives), return fd==65536 (FD_IDENT). 
 * This says that we can't open it, but should ident it.
 */


int  imager::open_dev(const char *friendly_name)
{
    /****************************************************************
     *** Check for ata%d or ide%d
     ****************************************************************/

    ata_dev = -1;
    sscanf(friendly_name,"ata%d",&ata_dev);	// try to find ata0
    if(ata_dev==-1){
	sscanf(friendly_name,"ide%d",&ata_dev);	// try to find ide0
    }

    if(ata_dev != -1){			// if we found the device
	/* Create the attach and detach commands */

	char dev0[64];			// space for the first channel
	char dev1[64];			// space for the second channel
	char *dev[2] = {dev0,dev1};

	make_ata_attach_commands(cmd_attach,cmd_detach,dev0,dev1,ata_dev);
	system(cmd_detach);	// make sure we are detached first

	int i;
	for(i=0;i<10;i++){
	    int delay = i*3;
	    printf("\nOpening special ATA Bus #%d...\n",ata_dev);
	    if(i>0){
		printf("Attempt %d out of %d.\n",i+1,10);
	    }
	    printf("# %s\n",cmd_attach);
	    system(cmd_attach);
	    
	    /* See if we found the device */
	    for(int d=0;d<2;d++){
		if(access(dev[d],F_OK)==0){
		    if(access(dev[d],R_OK)){
			// don't have permission to read it.
			// this is bad
			err(1,dev[d]);
		    }
		    if(delay){
			printf("Waiting %d second%s for %s to spin up...\n",
			       delay,delay==1?"":"s",infile);  
			isleep(delay);
		    }	
		    strcpy(infile,dev[d]); // we will try this one
		    int fd = open(infile,O_RDONLY);
		    if(fd>0){
			/* The device was successfully opened. */
			return fd;		// got it!
		    }
		    perror(infile);	       
		}
	    }
	    printf("Detaching device and trying again...\n");
	    printf("# %s\n",cmd_detach);
	    system(cmd_detach);
	    isleep(delay);
	}
	/* Been through too many times. Did we get a device?
	 * If so, just ident it...
	 */
	if(infile[0]){
	    imaging_failed = true;
	    return FD_IDENT;
	}
    }

    /****************************************************************
     *** Check for scsi%d
     *** In our testing with FreeBSD, there is no advantage to repeatedly
     *** attempting to attach or detach...
     ****************************************************************/
    if(sscanf(friendly_name,"scsi%d",&scsi_bus)==1){
	if(scsi_attach(infile,sizeof(infile),this)==0){
	    int fd = open(infile,O_RDONLY);
	    if(fd>0){
		return fd;
	    }
	    /* attach was successful but open failed. */
	    imaging_failed = true;
	    return FD_IDENT;
	}
    }
    return -1;
}





/*
 * main image loop.
 * if high_water_mark==0, then we do not know how many blocks the input
 * is; just read it byte-by-byte...
 */
void imager::image_loop(uint64 low_water_mark, // sector # to start
			uint64 high_water_mark, // sector # to end
			int direction, int readsectors,int error_mask)
{
    // buffer to store the data we read
    bufsize = readsectors*sector_size;
    buf = (unsigned char *)malloc(bufsize); 
    memset(buf,0,sizeof(buf));
    uint64 data_offset = 0;		// offset into output file
    bool valid_reverse_data = false;		// did we ever get valid data in the reverse direction?
    bool last_read_short = false;
    int reminder = 0;

    if(!buf) err(1,"malloc");

    /* Get the badflag that we'll be using */
    badflag = (unsigned char *)malloc(sector_size);
    if(af) memcpy(badflag,af_badflag(af),sector_size);
    else memset(badflag,0,sector_size);

    /* Loop as long as we have room, or until we get an EOF
     * (if high_water_mark is 0.)
     */
    imaging = true;
    while(low_water_mark < high_water_mark || high_water_mark==0){ 

	/* Figure out where to read and how how many sectors to read */
	uint64 snum;	// where we will be reading
	unsigned int sectors_to_read = readsectors;
	if(sectors_to_read > maxreadblocks && maxreadblocks>0){
	    sectors_to_read = maxreadblocks;
	}
	if(direction==1){	// going up
	    snum = low_water_mark;

	    /* If a high water mark is set, take it into account */
	    if(high_water_mark>0){
		unsigned int sectors_left = high_water_mark - snum;
		if(sectors_left < sectors_to_read){ 
		    sectors_to_read = sectors_left;
		}
	    }
	}
	else {
	    assert(high_water_mark != 0); // we can't go backwards if we don't know end
	    snum = high_water_mark - sectors_to_read;
	    if(snum<low_water_mark){
		snum = low_water_mark;
		sectors_to_read = high_water_mark - low_water_mark;
	    }
	}

	last_sector_read = snum;
	last_sectors_read = sectors_to_read;
	last_direction = direction;

	if (high_water_mark != 0){ 	// if we know where the top is...
	    data_offset = sector_size * snum; // where we want to start reading

	    if(data_offset != in_pos){	// eliminate unnecessary seeks
		lseek(in,data_offset,SEEK_SET);	// make sure we are at the right place; (ignore error)
		in_pos = data_offset;
	    }
	}

	status();			// tell the user what we are doing

	int bytes_to_read = sectors_to_read * sector_size;

	/* Fill the buffer that we are going to read with the bad flag */
	for(int i=0;i<bytes_to_read;i+=sector_size){
	    memcpy(buf+i,badflag,sector_size);
	}

	/* Now seek and read */

	int bytes_read    = 0;
	if(opt_debug==99){
	    bytes_read = -1; // simulate a read error
	} else {
	    if(opt_use_timers) read_timer.start();
	    bytes_read = read(in,buf,bytes_to_read);
	    if(opt_use_timers) read_timer.stop();
	}
	if(bytes_read>=0){
	    in_pos += bytes_read;	// update position
	}

	/* Note if we got valid data in the reverse direction */
	if((direction == -1) && (bytes_read>0)) valid_reverse_data = true;


	if(bytes_read == bytes_to_read){
	    /* Got a good read! */
	    total_sectors_read    += sectors_to_read;
	    total_bytes_read      += bytes_read;

	    /* Reset the error counters */
	    consecutive_read_errors = 0;
	    consecutive_read_error_regions = 0;
	    last_read_short = false;

	    /* Write the data! */
	    write_data(buf,data_offset,bytes_read);

	    if(direction==1){
		low_water_mark += sectors_to_read;
	    }
	    else {
		high_water_mark -= sectors_to_read;
	    }
	    continue;
	}

	/* Some kind of error... */

	/* If high water mark is 0,
	 * then just write out what we read and continue, because we don't know how many
	 * bytes we can read...
	 */
	if(high_water_mark==0 && bytes_read<=0){
	    break;	// end of pipe/file/whatever
	}

	/* If we are reading forward and we got an incomplete read, just live with it... */
	if(direction==1 && bytes_read>0){
	    total_bytes_read      += bytes_read; 
	    write_data(buf,data_offset,bytes_read);
	    data_offset += bytes_read;	// move along
	    low_water_mark += (bytes_read + reminder)/sector_size;
	    reminder = (bytes_read + reminder)%sector_size;
	    last_read_short = true;
	    continue;
	}

	/* Error handling follows. This code will automatically retry
	 * the same set of sectors retry_count and then switch direction.
	 */
	if(error_mask==0){
	    /* If errors on this attempted read exceed the threshold,
	     * just note how many bytes we were able to read and swap directions if necessary.
	     * If we have done that too many times in a row, then give up...
	     */
	    if(++consecutive_read_errors>retry_count){
		consecutive_read_errors=0; // reset the counter

		/* If we got an error, note it --- unless one of two conditions are true:
		 * we are going forwards and the last was a short read.
		 * we are going backwards and we have never gotten valid data going backwards.