// ****************************************************************************
//  Project:        GUYMAGER
// ****************************************************************************
//  Programmer:     Guy Voncken
//                  Police Grand-Ducale
//                  Service de Police Judiciaire
//                  Section Nouvelles Technologies
// ****************************************************************************
//  Module:         Multithreaded AFF (AAFF = Avanced AFF)
// ****************************************************************************

#include "common.h"
#include "compileinfo.h"

#include <netinet/in.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>
#include <zlib.h>

#include <QString>

#include "util.h"
#include "config.h"
#include "aaff.h"

// -----------------
//  AFF definitions
// -----------------

#define AFF_GID_LENGTH   16
#define AFF_SEGARG_U64   2  // Used as argument for segments that contain a 64 bit unsigned in the data field

#define AFF_HEADER                    "AFF10\r\n"
#define AFF_SEGMENT_HEADER_MAGIC      "AFF"
#define AFF_SEGMENT_FOOTER_MAGIC      "ATT"
#define AFF_BADSECTOR_HEADER          "BAD SECTOR"
#define AFF_FILE_TYPE                 "AFF"

#define AFF_SEGNAME_BADFLAG           "badflag"
#define AFF_SEGNAME_AFFLIB_VERSION    "afflib_version"
#define AFF_SEGNAME_FILETYPE          "aff_file_type"
#define AFF_SEGNAME_GID               "image_gid"
#define AFF_SEGNAME_SECTORS           "devicesectors"
#define AFF_SEGNAME_SECTORSIZE        "sectorsize"
#define AFF_SEGNAME_IMAGESIZE         "imagesize"
#define AFF_SEGNAME_PAGESIZE          "pagesize"
#define AFF_SEGNAME_BADSECTORS        "badsectors"
#define AFF_SEGNAME_MD5               "md5"
#define AFF_SEGNAME_SHA256            "sha256"
#define AFF_SEGNAME_DURATION          "acquisition_seconds"
#define AFF_SEGNAME_PAGE              "page"

#define AFF_PAGEFLAGS_UNCOMPRESSED    0x0000
#define AFF_PAGEFLAGS_COMPRESSED_ZLIB 0x0001
#define AFF_PAGEFLAGS_COMPRESSED_ZERO 0x0033    // Compressed, Zero and MaxCompression

typedef struct
{
   char         Magic[4];
   unsigned int NameLen;
   unsigned int DataLen;
   unsigned int Argument;          // Named "flags" in original aff source, named "arg" in afinfo output.
} t_AffSegmentHeader;

// Between header and footer lie the segment name and the data

typedef struct
{
   char         Magic[4];
   unsigned int SegmentLen;
} t_AffSegmentFooter;


// ------------------
//  Aaff definitions
// ------------------

typedef struct _t_Aaff
{
   FILE               *pFile;
   unsigned long long   PagesWritten;

   t_AffSegmentHeader   SegmentHeader;  // Optimisation: This header and this footer only need to be
   t_AffSegmentFooter   SegmentFooter;  // allocated and initialised once, and can be used again and again
} t_Aaff;

#define AAFF_MD5_LEN                16
#define AAFF_SHA256_LEN             32
#define AAFF_BADSECTORMARKER_MAXLEN 65536
unsigned char AaffBadSectorMarker[AAFF_BADSECTORMARKER_MAXLEN];

// ----------------
//  Error handling
// ----------------

#define CHK_FWRITE(Fn)                    \
   if ((Fn) != 1)                         \
      CHK (ERROR_AAFF_CANNOT_WRITE_FILE)


// -------------------
//  Utility functions
// -------------------

static bool AaffIsZero (unsigned char *pData, int DataLen)
{
   long long *pBuff    = (long long *)pData;
   const int   CmpSize = sizeof (long long);   // 64 bit operations for optimal performance on amd64 processors

   while (DataLen >= CmpSize)
   {
      if (*pBuff++)
        return false;
      DataLen -= CmpSize;
   }

   pData = (unsigned char *) pBuff;
   while (DataLen--)
   {
      if (*pData++)
         return false;
   }

   return true;
}

// -------------------
//  Segment functions
// -------------------

static APIRET AaffWriteSegment (t_pAaff pAaff, const char *pName, unsigned int Argument, const unsigned char *pData, unsigned int DataLen)
{
   int NameLen0 = strlen(pName);

   pAaff->SegmentHeader.NameLen    = htonl(NameLen0);
   pAaff->SegmentHeader.DataLen    = htonl(DataLen);
   pAaff->SegmentHeader.Argument   = htonl(Argument);
   pAaff->SegmentFooter.SegmentLen = htonl(sizeof(t_AffSegmentHeader) + sizeof(t_AffSegmentFooter) + NameLen0 + DataLen);

   CHK_FWRITE(fwrite (&pAaff->SegmentHeader, sizeof(t_AffSegmentHeader), 1, pAaff->pFile))
   CHK_FWRITE(fwrite (pName, NameLen0, 1, pAaff->pFile))
   if (pData && DataLen)
      CHK_FWRITE(fwrite (pData, DataLen, 1, pAaff->pFile))
   CHK_FWRITE(fwrite (&pAaff->SegmentFooter, sizeof(t_AffSegmentFooter), 1, pAaff->pFile))

   return NO_ERROR;
}

APIRET AaffWriteSegmentStr (t_pAaff pAaff, const char *pName, unsigned int Argument, const char *pStr)
{
   CHK (AaffWriteSegment (pAaff, pName, Argument, (const unsigned char *) pStr, strlen (pStr)))
   return NO_ERROR;
}

static APIRET AaffWriteSegmentArg (t_pAaff pAaff, const char *pName, unsigned int Argument)
{
   CHK (AaffWriteSegment (pAaff, pName, Argument, NULL, 0))
   return NO_ERROR;
}

static APIRET AaffWriteSegmentU64 (t_pAaff pAaff, const char *pName, unsigned long long Value)
{
   unsigned int Data[2];

   Data[0] = htonl ((unsigned int)(Value &  0xFFFFFFFF));
   Data[1] = htonl ((unsigned int)(Value >> 32));

   CHK (AaffWriteSegment (pAaff, pName, AFF_SEGARG_U64, (unsigned char *)&Data[0], sizeof(Data)))

   return NO_ERROR;
}

static APIRET AaffWriteSegmentGID (t_pAaff pAaff)
{
   unsigned char GID[AFF_GID_LENGTH];
   int           i;

   for (i=0; i<AFF_GID_LENGTH; i++)
      GID[i] = (unsigned char) random();

   CHK (AaffWriteSegment (pAaff, AFF_SEGNAME_GID, 0, GID, AFF_GID_LENGTH))

   return NO_ERROR;
}


// ---------------
//  API functions
// ---------------

APIRET AaffOpen (t_pAaff *ppAaff, const char *pFilename, unsigned long long DeviceSize, unsigned int SectorSize, unsigned int PageSize)
{
   t_pAaff pAaff;
   char     Buff[512];

   // Open file and intialise
   // -----------------------
   if (SectorSize > AAFF_BADSECTORMARKER_MAXLEN)
      CHK (ERROR_AAFF_SECTORSIZE_TOO_BIG)

   *ppAaff = NULL;
   pAaff = (t_pAaff) UTIL_MEM_ALLOC(sizeof(t_Aaff));
   if (pAaff == NULL)
      CHK (ERROR_AAFF_MEMALLOC_FAILED)

   pAaff->pFile = fopen64 (pFilename, "w");
   if (pAaff->pFile == NULL)
   {
      UTIL_MEM_FREE (pAaff);
      CHK (ERROR_AAFF_CANNOT_CREATE_FILE)
   }
   *ppAaff = pAaff;

   CHK_FWRITE(fwrite (AFF_HEADER, sizeof(AFF_HEADER), 1, pAaff->pFile))

   pAaff->PagesWritten = 0;
   memset (&pAaff->SegmentHeader, 0, sizeof (t_AffSegmentHeader));
   memset (&pAaff->SegmentFooter, 0, sizeof (t_AffSegmentFooter));
   strcpy (&pAaff->SegmentHeader.Magic[0], AFF_SEGMENT_HEADER_MAGIC);
   strcpy (&pAaff->SegmentFooter.Magic[0], AFF_SEGMENT_FOOTER_MAGIC);

   // Write standard segments
   // -----------------------
   snprintf (&Buff[0], sizeof(Buff), "aaff module of Guymager %s", pCompileInfoVersion);
   CHK (AaffWriteSegmentGID (pAaff))
   if CONFIG (AffMarkBadSectors)
      CHK (AaffWriteSegment (pAaff, AFF_SEGNAME_BADFLAG       , 0, AaffBadSectorMarker, SectorSize))
   CHK (AaffWriteSegmentStr (pAaff, AFF_SEGNAME_AFFLIB_VERSION, 0, &Buff[0]))
   CHK (AaffWriteSegmentStr (pAaff, AFF_SEGNAME_FILETYPE      , 0, AFF_FILE_TYPE))
   CHK (AaffWriteSegmentArg (pAaff, AFF_SEGNAME_PAGESIZE      , PageSize))
   CHK (AaffWriteSegmentArg (pAaff, AFF_SEGNAME_SECTORSIZE    , SectorSize))
   CHK (AaffWriteSegmentU64 (pAaff, AFF_SEGNAME_SECTORS       , DeviceSize / SectorSize))
   CHK (AaffWriteSegmentU64 (pAaff, AFF_SEGNAME_IMAGESIZE     , DeviceSize))

   return NO_ERROR;
}

APIRET AaffClose (t_pAaff pAaff, unsigned long long BadSectors, unsigned char *pMD5, unsigned char *pSHA256, int Duration)
{
   CHK (AaffWriteSegmentU64 (pAaff, AFF_SEGNAME_BADSECTORS  , BadSectors))
   CHK (AaffWriteSegment    (pAaff, AFF_SEGNAME_MD5         , 0, pMD5   , AAFF_MD5_LEN   ))
   CHK (AaffWriteSegment    (pAaff, AFF_SEGNAME_SHA256      , 0, pSHA256, AAFF_SHA256_LEN))
   CHK (AaffWriteSegmentArg (pAaff, AFF_SEGNAME_DURATION    , Duration))

   if (fflush (pAaff->pFile))
   {
      (void) fclose (pAaff->pFile);
      CHK (ERROR_AAFF_CANNOT_FLUSH_FILE)
   }

   if (fclose (pAaff->pFile))
      CHK (ERROR_AAFF_CANNOT_CLOSE_FILE)

   return NO_ERROR;
}

APIRET AaffPreprocess (t_pAaffPreprocess *ppPreprocess, unsigned char *pDataIn, unsigned int DataLenIn, unsigned char *pDataOut, unsigned int DataLenOut)
{
   t_pAaffPreprocess pPreProcess;
   int                rc;
   uLongf             LenOut;

   *ppPreprocess = NULL;
   pPreProcess = (t_pAaffPreprocess) UTIL_MEM_ALLOC(sizeof(t_AaffPreprocess));
   if (pPreProcess == NULL)
      CHK (ERROR_AAFF_MEMALLOC_FAILED)
   *ppPreprocess = pPreProcess;
   pPreProcess->Zero       = false;
   pPreProcess->Compressed = false;

   // Check if zero
   // -------------
   pPreProcess->Zero = AaffIsZero (pDataIn, DataLenIn);
   if (pPreProcess->Zero)
      return NO_ERROR;

   // Try compression
   // ---------------
   LenOut = DataLenOut;
   rc = compress2 ((Bytef *)pDataOut, &LenOut, (Bytef *)pDataIn, DataLenIn, CONFIG (AffCompression));
   pPreProcess->DataLenOut = LenOut;
   if (rc != Z_OK)
   {
      if (rc != Z_BUF_ERROR)    // Do not log this one (the destination buffer was too small for the compressed result)
         LOG_ERROR ("compress2 returned %d", rc)
      return NO_ERROR;
   }

   pPreProcess->Compressed = (LenOut < DataLenIn);

   return NO_ERROR;
}

APIRET AaffWrite (t_pAaff pAaff, t_pAaffPreprocess pPreprocess, unsigned char *pData, unsigned int DataLen)
{
   char SegmentName[64];

   snprintf (&SegmentName[0], sizeof(SegmentName), "%s%llu", AFF_SEGNAME_PAGE, pAaff->PagesWritten++);

   if (pPreprocess->Zero)
   {
      int Len = htonl(DataLen);
      CHK (AaffWriteSegment (pAaff, &SegmentName[0], AFF_PAGEFLAGS_COMPRESSED_ZERO, (unsigned char *) &Len, 4))
   }
   else if (pPreprocess->Compressed)
      CHK (AaffWriteSegment (pAaff, &SegmentName[0], AFF_PAGEFLAGS_COMPRESSED_ZLIB, pData, pPreprocess->DataLenOut))
   else
      CHK (AaffWriteSegment (pAaff, &SegmentName[0], AFF_PAGEFLAGS_UNCOMPRESSED   , pData, DataLen))

   UTIL_MEM_FREE (pPreprocess);

   return NO_ERROR;
}

APIRET AaffCopyBadSectorMarker (unsigned char *pBuffer, unsigned int Len)
{
   if (Len > AAFF_BADSECTORMARKER_MAXLEN)
      CHK (ERROR_AAFF_SECTORSIZE_TOO_BIG)

   memcpy (pBuffer, &AaffBadSectorMarker[0], Len);

   return NO_ERROR;
}

// -----------------------
//  Module initialisation
// -----------------------

APIRET AaffInit (void)
{
   int i;

   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_MEMALLOC_FAILED   ))
   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_CANNOT_CREATE_FILE))
   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_CANNOT_WRITE_FILE ))
   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_CANNOT_CLOSE_FILE ))
   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_CANNOT_FLUSH_FILE ))
   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_SECTORSIZE_TOO_BIG))
   CHK (TOOL_ERROR_REGISTER_CODE (ERROR_AAFF_COMPRESSION_FAILED))

   srandom (time (0));

   for (i=0; i<AAFF_BADSECTORMARKER_MAXLEN; i++)
      AaffBadSectorMarker[i] = (char)random();
   strcpy ((char *)AaffBadSectorMarker, AFF_BADSECTOR_HEADER);

   return NO_ERROR;
}

APIRET AaffDeInit (void)
{
   return NO_ERROR;
}