Skip to content
Snippets Groups Projects
Commit 38041479 authored by Ann Priestman's avatar Ann Priestman
Browse files

Implemented image writer add

parent 8964f703
Branches
Tags
No related merge requests found
Hide whitespace changes
Inline Side-by-side
tsk/img/img_writer.cpp
+ 274
− 20
View file @ 38041479
......@@ -29,21 +29,267 @@
#define VHD_FOOTER_LENGTH 0x200
#define VHD_DISK_HEADER_LENGTH 0x400
static TSK_RETVAL_ENUM tsk_img_writer_add(TSK_IMG_WRITER* img_writer, TSK_OFF_T addr, char *buffer, size_t len) {
TSK_RETVAL_ENUM writeFooter(TSK_IMG_WRITER* writer);
bool getBit(unsigned char * buffer, TSK_OFF_T index) {
unsigned char b = buffer[index / 8];
return (b >> (7 - (index % 8)) & 0x01) == 1;
}
void setBit(unsigned char * buffer, TSK_OFF_T index, bool val) {
unsigned char b = buffer[index / 8];
unsigned char mask = 0xff ^ (1 << (7 - (index % 8)));
b = (b & mask) | (val << (7 - (index % 8)));
buffer[index / 8] = b;
}
TSK_RETVAL_ENUM seekToOffset(TSK_IMG_WRITER * writer, TSK_OFF_T offset) {
LARGE_INTEGER li;
li.QuadPart = offset;
li.LowPart = SetFilePointer(writer->outputFileHandle, li.LowPart,
&li.HighPart, FILE_BEGIN);
if ((li.LowPart == INVALID_SET_FILE_POINTER) &&
(GetLastError() != NO_ERROR)) {
int lastError = (int)GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_SEEK);
tsk_error_set_errstr("img_writer::seekToOffset: offset %" PRIuOFF " seek - %d",
offset,
lastError);
return TSK_ERR;
}
return TSK_OK;
}
TSK_RETVAL_ENUM seekAhead(TSK_IMG_WRITER * writer, TSK_OFF_T dist) {
LARGE_INTEGER li;
li.QuadPart = dist;
li.LowPart = SetFilePointer(writer->outputFileHandle, li.LowPart,
&li.HighPart, FILE_CURRENT);
if ((li.LowPart == INVALID_SET_FILE_POINTER) &&
(GetLastError() != NO_ERROR)) {
int lastError = (int)GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_SEEK);
tsk_error_set_errstr("img_writer::seekAhead: offset %" PRIuOFF " seek - %d",
dist,
lastError);
return TSK_ERR;
}
return TSK_OK;
}
void checkIfBlockIsFinished(TSK_IMG_WRITER* writer, TSK_OFF_T blockNum) {
/* This may not ever mark the last block as finished - need to work on that */
unsigned char * sectBitmap = writer->blockToSectorBitmap[blockNum];
for (int i = 0; i < writer->sectorsPerBlock; i++) {
if (false == getBit(sectBitmap, i)) {
/* At least one sector has not been written */
return;
}
}
/* Mark the block as finished and free the memory for its sector bitmap */
writer->blockStatus[blockNum] = IMG_WRITER_BLOCK_STATUS_FINISHED;
if (writer->blockToSectorBitmap[blockNum] != NULL) {
free(writer->blockToSectorBitmap[blockNum]);
}
}
TSK_RETVAL_ENUM addToExistingBlock(TSK_IMG_WRITER* writer, TSK_OFF_T addr, char *buffer,
size_t len, TSK_OFF_T blockNum) {
tsk_fprintf(stdout, "Block %d adding to existing block starting at rel offset 0x%llx\n",
blockNum, addr % writer->blockSize);
/* Seek to where this buffer should start in the image */
if (TSK_OK != seekToOffset(writer, writer->blockToOffset[blockNum] + writer->sectorBitmapLength +
(addr % writer->blockSize))) {
return TSK_ERR;
}
/* Copy each sector that isn't already there */
for (TSK_OFF_T inputOffset = 0; inputOffset < len; inputOffset += VHD_SECTOR_SIZE) {
uint32_t currentSector = (addr % writer->blockSize + inputOffset) / VHD_SECTOR_SIZE;
if (getBit(writer->blockToSectorBitmap[blockNum], currentSector)) {
if (TSK_OK != seekAhead(writer, VHD_SECTOR_SIZE)) {
return TSK_ERR;
}
}
else {
DWORD bytesWritten;
if (FALSE == WriteFile(writer->outputFileHandle, &(buffer[inputOffset]), VHD_SECTOR_SIZE,
&bytesWritten, NULL)) {
int lastError = GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_WRITE);
tsk_error_set_errstr("addToExistingBlock: error writing sector",
lastError);
return TSK_ERR;
}
setBit(writer->blockToSectorBitmap[blockNum], currentSector, true);
}
}
/* Update the sector bitmap */
if (TSK_OK != seekToOffset(writer, writer->blockToOffset[blockNum])) {
return TSK_ERR;
}
DWORD bytesWritten;
if (FALSE == WriteFile(writer->outputFileHandle, writer->blockToSectorBitmap[blockNum],
writer->sectorBitmapArrayLength, &bytesWritten, NULL)) {
int lastError = GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_WRITE);
tsk_error_set_errstr("addToExistingBlock: error writing sector",
lastError);
return TSK_ERR;
}
return TSK_OK;
}
TSK_RETVAL_ENUM addNewBlock(TSK_IMG_WRITER* writer, TSK_OFF_T addr, char *buffer, size_t len, TSK_OFF_T blockNum) {
writer->blockStatus[blockNum] = IMG_WRITER_BLOCK_STATUS_ALLOC;
writer->blockToOffset[blockNum] = writer->nextDataOffset;
char * fullBuffer = (char *)tsk_malloc(writer->blockSize * sizeof(char));
char * sectorBitmap = (char *)tsk_malloc(writer->sectorBitmapLength * sizeof(char));
unsigned char * completedSectors = (unsigned char *)tsk_malloc(((writer->sectorBitmapArrayLength) * sizeof(char)));
/* Create the full new block and record the sectors written */
TSK_OFF_T startingOffset = addr % writer->blockSize;
for (size_t i = 0; i < len; i++) {
if (((startingOffset + i) % VHD_SECTOR_SIZE) == 0) {
TSK_OFF_T currentSector = (startingOffset + i) / VHD_SECTOR_SIZE;
setBit(completedSectors, currentSector, true);
sectorBitmap[currentSector / 8] = completedSectors[currentSector / 8];
}
fullBuffer[startingOffset + i] = buffer[i];
}
writer->blockToSectorBitmap[blockNum] = completedSectors;
/* Prepare the new block offset */
TSK_OFF_T nextDataOffsetSector = writer->nextDataOffset / VHD_SECTOR_SIZE;
unsigned char newBlockOffset[4];
newBlockOffset[0] = (nextDataOffsetSector >> 24) & 0xff;
newBlockOffset[1] = (nextDataOffsetSector >> 16) & 0xff;
newBlockOffset[2] = (nextDataOffsetSector >> 8) & 0xff;
newBlockOffset[3] = nextDataOffsetSector & 0xff;
/* Write the new offset to the BAT */
if (TSK_OK != seekToOffset(writer, writer->batOffset + 4 * blockNum)) {
return TSK_ERR;
}
DWORD bytesWritten;
if (FALSE == WriteFile(writer->outputFileHandle, newBlockOffset, 4, &bytesWritten, NULL)) {
int lastError = GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_WRITE);
tsk_error_set_errstr("addNewBlock: error writing BAT entry",
lastError);
return TSK_ERR;
}
/* Write the sector bitmap and the data */
if (TSK_OK != seekToOffset(writer, writer->nextDataOffset)) {
return TSK_ERR;
}
if (FALSE == WriteFile(writer->outputFileHandle, sectorBitmap, writer->sectorBitmapLength, &bytesWritten, NULL)) {
int lastError = GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_WRITE);
tsk_error_set_errstr("addNewBlock: error writing sector bitmap",
lastError);
return TSK_ERR;
}
if (FALSE == WriteFile(writer->outputFileHandle, fullBuffer, writer->blockSize, &bytesWritten, NULL)) {
int lastError = GetLastError();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_WRITE);
tsk_error_set_errstr("addNewBlock: error writing block data",
lastError);
return TSK_ERR;
}
writer->nextDataOffset += writer->sectorBitmapLength + writer->blockSize;
/* Always add the footer on to make it a valid VHD */
writeFooter(writer);
free(fullBuffer);
free(sectorBitmap);
return TSK_OK;
}
TSK_RETVAL_ENUM addBlock(TSK_IMG_WRITER* writer, TSK_OFF_T addr, char *buffer, size_t len) {
TSK_OFF_T blockNum = addr / writer->blockSize;
if (writer->blockStatus[blockNum] == IMG_WRITER_BLOCK_STATUS_FINISHED){
return TSK_OK;
}
fflush(stderr);
if (writer->blockStatus[blockNum] == IMG_WRITER_BLOCK_STATUS_ALLOC) {
addToExistingBlock(writer, addr, buffer, len, blockNum);
}
else {
addNewBlock(writer, addr, buffer, len, blockNum);
}
/* Check whether the block is now done */
checkIfBlockIsFinished(writer, blockNum);
return TSK_OK;
}
static TSK_RETVAL_ENUM tsk_img_writer_add(TSK_IMG_WRITER* writer, TSK_OFF_T addr, char *buffer, size_t len) {
#ifndef TSK_WIN32
return TSK_ERR;
return TSK_ERR;
#else
if (tsk_verbose) {
tsk_fprintf(stderr,
"tsk_img_writer_add: Adding data at offset: %"
PRIuOFF " len: %" PRIuOFF "\n", addr,
(TSK_OFF_T)len);
}
return TSK_OK;
if (tsk_verbose) {
tsk_fprintf(stderr,
"tsk_img_writer_add: Adding data at offset: %"
PRIuOFF " len: %" PRIuOFF "\n", addr,
(TSK_OFF_T)len);
}
/* This should never happen, but best to check */
if (addr % VHD_SECTOR_SIZE != 0) {
return TSK_ERR;
}
if (addr % writer->blockSize == 0) {
/* The buffer matches up with the blocks */
return addBlock(writer, addr, buffer, len);
}
else {
/* The buffer spans two blocks */
TSK_OFF_T firstPartLength = writer->blockSize - (addr % writer->blockSize);
addBlock(writer, addr, buffer, firstPartLength);
if (addr + firstPartLength < writer->imageSize) {
addBlock(writer, addr + firstPartLength, buffer + firstPartLength, addr % writer->blockSize);
}
}
return TSK_OK;
#endif
}
TSK_RETVAL_ENUM writeFooter(TSK_IMG_WRITER* writer);
static TSK_RETVAL_ENUM tsk_img_writer_close(TSK_IMG_WRITER* img_writer) {
#ifndef TSK_WIN32
return TSK_ERR;
......@@ -53,10 +299,10 @@ static TSK_RETVAL_ENUM tsk_img_writer_close(TSK_IMG_WRITER* img_writer) {
"tsk_img_writer_close: Closing image writer");
}
if (img_writer->outputFileHandle != 0) {
TSK_RETVAL_ENUM retval = writeFooter(img_writer);
CloseHandle(img_writer->outputFileHandle);
img_writer->outputFileHandle = 0;
return retval;
/* TODO: Free the data */
}
return TSK_OK;
......@@ -112,14 +358,12 @@ uint32_t generateChecksum(unsigned char * buffer, int len) {
/*
* Write the footer (which is also the first sector) to the file.
* Will write to the current file position.
* Save it so we only have to generate it once.
*/
TSK_RETVAL_ENUM writeFooter(TSK_IMG_WRITER* writer) {
if (writer->footer == NULL) {
writer->footer = (unsigned char *)malloc(VHD_FOOTER_LENGTH * sizeof(unsigned char));
for (int i = 0; i < VHD_FOOTER_LENGTH; i++) {
writer->footer[i] = 0;
}
writer->footer = (unsigned char *)tsk_malloc(VHD_FOOTER_LENGTH * sizeof(unsigned char));
/* First calculate geometry values */
uint32_t cylinders;
......@@ -213,7 +457,7 @@ TSK_RETVAL_ENUM writeDynamicDiskHeader(TSK_IMG_WRITER * writer) {
tsk_error_reset();
tsk_error_reset();
tsk_error_set_errno(TSK_ERR_IMG_WRITE);
tsk_error_set_errstr("writeFooter: error writing VHD footer",
tsk_error_set_errstr("writeFooter: error writing VHD header",
lastError);
return TSK_ERR;
}
......@@ -230,9 +474,6 @@ TSK_RETVAL_ENUM tsk_img_writer_create(TSK_IMG_INFO * img_info, const TSK_TCHAR *
#ifndef TSK_WIN32
return TSK_ERR;
#else
tsk_fprintf(stdout,
"tsk_img_writer_create: Creating image writer in directory %" PRIttocTSK" with basename %" PRIttocTSK"\n",
directory, basename);
if (tsk_verbose) {
tsk_fprintf(stderr,
"tsk_img_writer_create: Creating image writer in directory %" PRIttocTSK" with basename %" PRIttocTSK"\n",
......@@ -267,6 +508,14 @@ TSK_RETVAL_ENUM tsk_img_writer_create(TSK_IMG_INFO * img_info, const TSK_TCHAR *
writer->totalBlocks++;
}
writer->sectorsPerBlock = writer->blockSize / VHD_SECTOR_SIZE;
writer->sectorBitmapArrayLength = writer->sectorsPerBlock / 8;
if (writer->sectorBitmapArrayLength % 8 != 0) {
writer->sectorBitmapArrayLength++;
}
writer->sectorBitmapLength = writer->sectorBitmapArrayLength;
if (0 != writer->sectorBitmapLength % VHD_SECTOR_SIZE) {
writer->sectorBitmapLength += VHD_SECTOR_SIZE - (writer->sectorBitmapLength % VHD_SECTOR_SIZE);
}
/* Set up the output file */
size_t len = TSTRLEN(directory) + TSTRLEN(basename) + 10;
......@@ -329,6 +578,11 @@ TSK_RETVAL_ENUM tsk_img_writer_create(TSK_IMG_INFO * img_info, const TSK_TCHAR *
/* Offset for the first data block - 0x600 bytes for the two headers plus the BAT length*/
writer->nextDataOffset = 0x600 + batLengthOnDisk;
/* Initialze all the bookkeeping arrays */
writer->blockStatus = (IMG_WRITER_BLOCK_STATUS_ENUM*)tsk_malloc(writer->totalBlocks * sizeof(IMG_WRITER_BLOCK_STATUS_ENUM));
writer->blockToOffset = (TSK_OFF_T*)tsk_malloc(writer->totalBlocks * sizeof(TSK_OFF_T));
writer->blockToSectorBitmap = (unsigned char **)tsk_malloc(writer->totalBlocks * sizeof(unsigned char *));
return TSK_OK;
#endif
......
tsk/img/img_writer.h
+ 12
− 0
View file @ 38041479
......@@ -23,6 +23,13 @@ extern "C" {
TSK_RETVAL_ENUM tsk_img_writer_create(TSK_IMG_INFO* img_info, const TSK_TCHAR * directory,
const TSK_TCHAR * basename);
enum IMG_WRITER_BLOCK_STATUS_ENUM {
IMG_WRITER_BLOCK_STATUS_UNALLOC = 0,
IMG_WRITER_BLOCK_STATUS_ALLOC = 1,
IMG_WRITER_BLOCK_STATUS_FINISHED = 2
};
typedef enum IMG_WRITER_BLOCK_STATUS_ENUM IMG_WRITER_BLOCK_STATUS_ENUM;
typedef struct TSK_IMG_WRITER TSK_IMG_WRITER;
struct TSK_IMG_WRITER {
TSK_IMG_INFO img_info;
......@@ -37,10 +44,15 @@ extern "C" {
TSK_OFF_T imageSize;
uint32_t totalBlocks;
uint32_t sectorBitmapLength;
uint32_t sectorBitmapArrayLength;
uint32_t sectorsPerBlock;
TSK_OFF_T batOffset;
TSK_OFF_T nextDataOffset;
IMG_WRITER_BLOCK_STATUS_ENUM* blockStatus;
TSK_OFF_T* blockToOffset;
unsigned char ** blockToSectorBitmap;
TSK_RETVAL_ENUM(*add)(TSK_IMG_WRITER* img_writer, TSK_OFF_T addr, char *buffer, size_t len);
TSK_RETVAL_ENUM(*close)(TSK_IMG_WRITER* img_writer);
TSK_RETVAL_ENUM(*finish_image)(TSK_IMG_WRITER* img_writer);
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment