ssn-installer/src/main.c

#include <getopt.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

//Import our code
#include "decrypt.h"
#include "fileReader.h"
#include "inflate.h"
#include "utils/min.h"

#define BUFFER_SIZE 512UL * 1024UL //512 KiB
#define ENCRYPTION_HEADER_LENGTH 12UL
#define LOCAL_FILE_HEADER_MAGIC (uint32_t)0x04034b50

uint16_t getUint16(uint8_t* buffer) {
  return (uint16_t)buffer[0] | \
    (uint16_t)buffer[1] << 8;
}
uint32_t getUint32(uint8_t* buffer) {
  return (uint32_t)buffer[0] | \
    (uint32_t)buffer[1] << 8 | \
    (uint32_t)buffer[2] << 16  | \
    (uint32_t)buffer[3] << 24;
}

struct arguments {
  /** Path to the disk file that contains the start of the file we want to extract.
   * The file may stretch across multiple disks though.
   */
  char* diskName;
  /** Offset into the disk where the file starts. */
  unsigned long diskOffset;
  /** Size of the file stored in the disk. */
  unsigned long fileSize;
  //Decryption keys
  bool isEncrypted;
  //For xdelta3, the location of the old file
  char* prevFile;
};

static struct option long_options[] = {
  {"disk", required_argument, 0, 'd'},
  {"offset", required_argument, 0, 'o'},
  {"size", required_argument, 0, 's'},
  {"keys", required_argument, 0, 'k'},
  {"prev", required_argument, 0, 'p'},
};

//Stores current state from command line arguments, initialized to zero
struct arguments state = {};

int main(int argc, char *argv[]) {
  //Parse command line arguments

  int requiredOptions = 0;

  while (1) {
    //in this variable, getopt_long stores the current position in the command line args array
    int option_index = 0;

    int curOption = getopt_long(argc, argv, "d:o:s:k:p:", long_options, &option_index);

    //end of command line arguments reached
    if (curOption == -1) {
      break;
    }

    switch (curOption) {
      case 'd': //disk name
        state.diskName = optarg;
        requiredOptions |= 1;
        break;
      case 'o': //offset
        state.diskOffset = atol(optarg);
        requiredOptions |= 2;
        break;
      case 's': //size
        state.fileSize = atol(optarg);
        requiredOptions |= 4;
        break;
      case 'k': { //decryption keys
        //TODO: parse from optarg
        uint32_t key0 = atoi(argv[8]);
        uint32_t key1 = atoi(argv[9]);
        uint32_t key2 = atoi(argv[10]);
        //Initialize decryption (pass decryption keys)
        initDecryptor(key0, key1, key2);
        state.isEncrypted = true;
        break;
      }
      case 'p': //prev file for xdelta3
        //TODO
        break;
      default:
        fprintf(stderr, "Unknown option '%c'.", (char)curOption);
        exit(1);
    }
  }

  if (requiredOptions != 7) {
    fprintf(stderr, "Missing arguments, received %i.", requiredOptions);
    exit(1);
  }

  //-------------------------------------------------

  uint8_t* compressedChunk = malloc(BUFFER_SIZE);
  if (compressedChunk == NULL) {
    fprintf(stderr, "Could not allocate %lu bytes for compressed buffer.\n", BUFFER_SIZE);
    exit(1);
  }

  uint8_t* uncompressedChunk = malloc(BUFFER_SIZE);
  if (uncompressedChunk == NULL) {
    fprintf(stderr, "Could not allocate %lu bytes for uncompressed buffer.\n", BUFFER_SIZE);
    exit(1);
  }
  memset(uncompressedChunk, (uint8_t)0, BUFFER_SIZE);

  //-------------------------------------------------

  //Initialize file reader
  initFileReader(state.diskName, state.diskOffset);

  //Skip local file header (30 bytes + additional length)
  getBytes(compressedChunk, 30UL);
  //Check that header is correct
  const uint32_t magic = getUint32(compressedChunk);
  if (magic != LOCAL_FILE_HEADER_MAGIC) {
    fprintf(stderr, "Wrong magic in local file header, expected %#010x but found %#010x.", LOCAL_FILE_HEADER_MAGIC, magic);
    exit(1);
  }
  //Read additional length
  const unsigned long additionalLength = getUint16(compressedChunk + 26) + getUint16(compressedChunk + 28);
  if (additionalLength > 0UL) {
    getBytes(NULL, additionalLength);
  }

  //If file is encrypted, skip 12-byte encryption header
  if (state.isEncrypted) {
    getBytes(compressedChunk, ENCRYPTION_HEADER_LENGTH);
    decrypt(compressedChunk, ENCRYPTION_HEADER_LENGTH);
  }

  //-------------------------------------------------

  struct InflateOutput inflateResult;
  inflateInit(compressedChunk, uncompressedChunk, BUFFER_SIZE);

  //Read actual file
  unsigned long remainingBytes = state.fileSize;
  bool needToRead = true;
  bool hasReachedEnd = false;
  unsigned long chunkSize;
  unsigned long uncompressedPosition = 0UL;
  while (remainingBytes > 0 || !hasReachedEnd) {
    if (needToRead) {
      chunkSize = min(BUFFER_SIZE - uncompressedPosition, remainingBytes);
      getBytes(compressedChunk, chunkSize);
      remainingBytes -= chunkSize;

      //Decrypt file if it is encrypted
      if (state.isEncrypted) {
        //TODO: For highest performance, we need to move if condition outside of while loop
        decrypt(compressedChunk, chunkSize);
      }
    }

    //Decompress file
    //bytes are contained in uncompressedChunk from [0, inflateResult.numBytesWrittenToOutput - 1]
    inflateResult = inflateInflate(needToRead ? chunkSize : 0, remainingBytes > 0);
    needToRead = inflateResult.needMoreInput;
    hasReachedEnd = inflateResult.hasReachedEnd;

    //important: we must not modify uncompressedChunk since miniz may use it as dictionary and read from it during the next invocation of inflateInflate()

    //Write to stdout
    write(1, uncompressedChunk + uncompressedPosition, inflateResult.numBytesWrittenToOutput);
    uncompressedPosition += inflateResult.numBytesWrittenToOutput;
    while (uncompressedPosition >= BUFFER_SIZE) {
      uncompressedPosition -= BUFFER_SIZE;
    }

    //Optionally perform xdelta3
    if (state.prevFile) {
      //TODO
    }
  }

  //release memory
  free(compressedChunk);
  free(uncompressedChunk);

  return 0;
}