/*
 * Copyright (c) 2008-2015 Freescale Semiconductor, Inc.
 * Copyright 2016-2018 NXP
 * All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include "stdafx.h"
#include <iostream>
#include <fstream>
#include <sstream>
#include <stdlib.h>
#include <stdexcept>
#include <stdio.h>
#include "options.h"
#include "EncoreBootImage.h"
#include "smart_ptr.h"
#include "Logging.h"
#include "format_string.h"
#include "Blob.h"
#include "Random.h"
#include "rijndael.h"

using namespace elftosb;

//! Size in bytes of the unencrypted group private key.
#define GPK_LENGTH (40)

//! Size in bytes of the encrypted output data. This size must be modulo 16, the chunk size for the
//! AES-128 crypto algorithm. The group private key is inserted at offset 16.
#define OUTPUT_DATA_LENGTH (64)

//! Position in the output data of the first byte of the group private key.
#define OUTPUT_DATA_GPK_OFFSET (16)

//! The tool's name.
const char k_toolName[] = "encryptgpk";

//! Current version number for the tool.
const char k_version[] = "1.0.2";

//! Copyright string.
const char k_copyright[] = "Copyright (c) 2008 Freescale Semiconductor. All rights reserved.";

//! Default output array name.
const char k_defaultArrayName[] = "_endDisplay";

//! Definition of command line options.
static const char *k_optionsDefinition[] = { "?|help",   "v|version", "k:key <file>", "z|zero-key", "o:output",
                                             "p:prefix", "a:array",   "d|debug",      "q|quiet",    "V|verbose",
                                             NULL };

//! Help string.
const char k_usageText[] =
    "\nOptions:\n\
  -?/--help                    Show this help\n\
  -v/--version                 Display tool version\n\
  -k/--key <file>              Add OTP key used for decryption\n\
  -z/--zero-key                Add default key of all zeroes\n\
  -o/--output <file>           Write output to this file\n\
  -p/--prefix <prefix>         Set the output array prefix\n\
  -a/--array <name>            Specify the output array name\n\
  -d/--debug                   Enable debug output\n\
  -q/--quiet                   Output only warnings and errors\n\
  -V/--verbose                 Print extra detailed log information\n\n";

//! Init vector used for CBC encrypting the output data.
static const uint8_t kInitVector[] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 };

//! An array of strings.
typedef std::vector<std::string> string_vector_t;

// prototypes
int main(int argc, char *argv[], char *envp[]);

/*!
 * \brief Class that encapsulates the sbtool interface.
 *
 * A single global logger instance is created during object construction. It is
 * never freed because we need it up to the last possible minute, when an
 * exception could be thrown.
 */
class encryptgpk
{
protected:
    int m_argc;                       //!< Number of command line arguments.
    char **m_argv;                    //!< String value for each command line argument.
    StdoutLogger *m_logger;           //!< Singleton logger instance.
    string_vector_t m_keyFilePaths;   //!< Paths to OTP key files.
    string_vector_t m_positionalArgs; //!< Arguments coming after explicit options.
    bool m_isVerbose;                 //!< Whether the verbose flag was turned on.
    bool m_useDefaultKey;             //!< Include a default (zero) crypto key.
    std::string m_outputPath;         //!< Path to output file.
    std::string m_gpkPath;            //!< Path to input group private key file.
    std::string m_outputPrefix;       //!< Prefix to the output array.
    std::string m_arrayName;          //!< Output array's name.

public:
    /*!
     * Constructor.
     *
     * Creates the singleton logger instance.
     */
    encryptgpk(int argc, char *argv[])
        : m_argc(argc)
        , m_argv(argv)
        , m_logger(0)
        , m_keyFilePaths()
        , m_positionalArgs()
        , m_isVerbose(false)
        , m_useDefaultKey(false)
        , m_outputPath()
        , m_gpkPath()
        , m_outputPrefix()
        , m_arrayName(k_defaultArrayName)
    {
        // create logger instance
        m_logger = new StdoutLogger();
        m_logger->setFilterLevel(Logger::INFO);
        Log::setLogger(m_logger);
    }

    /*!
     * Destructor.
     */
    ~encryptgpk() {}
    /*!
     * Reads the command line options passed into the constructor.
     *
     * This method can return a return code to its caller, which will cause the
     * tool to exit immediately with that return code value. Normally, though, it
     * will return -1 to signal that the tool should continue to execute and
     * all options were processed successfully.
     *
     * The Options class is used to parse command line options. See
     * #k_optionsDefinition for the list of options and #k_usageText for the
     * descriptive help for each option.
     *
     * \retval -1 The options were processed successfully. Let the tool run normally.
     * \return A zero or positive result is a return code value that should be
     *		returned from the tool as it exits immediately.
     */
    int processOptions()
    {
        Options options(*m_argv, k_optionsDefinition);
        OptArgvIter iter(--m_argc, ++m_argv);

        // process command line options
        int optchar;
        const char *optarg;
        while (optchar = options(iter, optarg))
        {
            switch (optchar)
            {
                case '?':
                    printUsage(options);
                    return 0;

                case 'v':
                    printf("%s %s\n%s\n", k_toolName, k_version, k_copyright);
                    return 0;

                case 'k':
                    m_keyFilePaths.push_back(optarg);
                    break;

                case 'z':
                    m_useDefaultKey = true;
                    break;

                case 'o':
                    m_outputPath = optarg;
                    break;

                case 'p':
                    m_outputPrefix = optarg;
                    break;

                case 'a':
                    m_arrayName = optarg;
                    break;

                case 'd':
                    Log::getLogger()->setFilterLevel(Logger::DEBUG);
                    break;

                case 'q':
                    Log::getLogger()->setFilterLevel(Logger::WARNING);
                    break;

                case 'V':
                    m_isVerbose = true;
                    break;

                default:
                    Log::log(Logger::ERROR, "error: unrecognized option\n\n");
                    printUsage(options);
                    return 1;
            }
        }

        // handle positional args
        if (iter.index() < m_argc)
        {
            //			Log::SetOutputLevel leveler(Logger::DEBUG);
            //			Log::log("positional args:\n");
            int i;
            for (i = iter.index(); i < m_argc; ++i)
            {
                //				Log::log("%d: %s\n", i - iter.index(), m_argv[i]);
                m_positionalArgs.push_back(m_argv[i]);
            }
        }

        // all is well
        return -1;
    }

    /*!
     * Prints help for the tool.
     */
    void printUsage(Options &options)
    {
        options.usage(std::cout, "gpk-file");
        printf(k_usageText, k_toolName);
    }

    /*!
     * Core of the tool. Calls processOptions() to handle command line options
     * before performing the real work the tool does.
     */
    int run()
    {
        try
        {
            // read command line options
            int result;
            if ((result = processOptions()) != -1)
            {
                return result;
            }

            // set verbose logging
            setVerboseLogging();

            // make sure a file was provided
            if (m_positionalArgs.size() < 1)
            {
                throw std::runtime_error("no input file path was provided");
            }

            // Make sure at least one key was specified.
            if (m_keyFilePaths.size() == 0 && m_useDefaultKey == false)
            {
                throw std::runtime_error("no crypto key was specified");
            }

            // Do the work.
            generateOutput();
        }
        catch (std::exception &e)
        {
            Log::log(Logger::ERROR, "error: %s\n", e.what());
            return 1;
        }
        catch (...)
        {
            Log::log(Logger::ERROR, "error: unexpected exception\n");
            return 1;
        }

        return 0;
    }

    /*!
     * \brief Builds the output data blob, encrypts it, and writes it to the output file.
     */
    void generateOutput()
    {
        // Create the output data blob and set it to the correct size.
        Blob data;
        data.setLength(OUTPUT_DATA_LENGTH);

        // Fill it with random values.
        RandomNumberGenerator rng;
        rng.generateBlock(data.getData(), OUTPUT_DATA_LENGTH);

        // Read the GPK and overlay it into the output data.
        // The first positional arg is the GPK file path.
        Blob gpk = readGPK(m_positionalArgs[0]);
        memcpy(data.getData() + OUTPUT_DATA_GPK_OFFSET, gpk.getData(), GPK_LENGTH);

        // This is the key object for our crypto key.
        AESKey<128> cryptoKey = readKeyFile();

        // Read the key file.
        // Encrypt the output data block.
        Rijndael cipher;
        cipher.init(Rijndael::CBC, Rijndael::Encrypt, cryptoKey, Rijndael::Key16Bytes, (uint8_t *)&kInitVector);
        cipher.blockEncrypt(data.getData(), OUTPUT_DATA_LENGTH * 8, data.getData());

        // Open the output file.
        std::ofstream outputStream(m_outputPath.c_str(), std::ios_base::out | std::ios_base::trunc);
        if (!outputStream.is_open())
        {
            throw std::runtime_error(format_string("could not open output file %s", m_outputPath.c_str()));
        }

        writeCArray(outputStream, data);
    }

    /*!
     * \brief Reads the group private key binary data.
     */
    Blob readGPK(std::string &path)
    {
        std::ifstream stream(path.c_str(), std::ios_base::in | std::ios_base::binary);
        if (!stream.is_open())
        {
            throw std::runtime_error("could not open group private key file");
        }

        Blob gpk;
        gpk.setLength(GPK_LENGTH);

        stream.read((char *)gpk.getData(), GPK_LENGTH);

        return gpk;
    }

    /*!
     * \brief Returns a key object based on the user's specified key.
     */
    AESKey<128> readKeyFile()
    {
        if (m_keyFilePaths.size() > 0)
        {
            // Open the key file.
            std::string &keyPath = m_keyFilePaths[0];
            std::ifstream keyStream(keyPath.c_str(), std::ios_base::in);
            if (!keyStream.is_open())
            {
                throw std::runtime_error(format_string("unable to read key file %s\n", keyPath.c_str()));
            }
            keyStream.seekg(0);

            // Read the first key in the file.
            AESKey<128> key(keyStream);
            return key;
        }

        // Otherwise, create a zero key and return it.
        AESKey<128> defaultKey;
        return defaultKey;
    }

    /*!
     * \brief Writes the given data blob as an array in a C source file.
     */
    void writeCArray(std::ofstream &stream, const Blob &data)
    {
        const uint8_t *dataPtr = data.getData();
        unsigned length = data.getLength();

        // Write first line.
        std::string text = format_string("%s%sunsigned char %s[%d] = {", m_outputPrefix.c_str(),
                                         m_outputPrefix.size() > 0 ? " " : "", m_arrayName.c_str(), length);
        stream.write(text.c_str(), text.size());

        // Write each word of the array.
        unsigned i = 0;
        while (i < length)
        {
            // Insert a comma at the end of the previous line unless this is the first word we're outputting.
            text = format_string("%s\n    0x%02x", i == 0 ? "" : ",", (*dataPtr++) & 0xff);
            stream.write(text.c_str(), text.size());

            i++;
        }

        // Write last line, terminating the array.
        text = "\n};\n\n";
        stream.write(text.c_str(), text.size());
    }

    /*!
     * \brief Turns on verbose logging.
     */
    void setVerboseLogging()
    {
        if (m_isVerbose)
        {
            // verbose only affects the INFO and DEBUG filter levels
            // if the user has selected quiet mode, it overrides verbose
            switch (Log::getLogger()->getFilterLevel())
            {
                case Logger::INFO:
                    Log::getLogger()->setFilterLevel(Logger::INFO2);
                    break;
                case Logger::DEBUG:
                    Log::getLogger()->setFilterLevel(Logger::DEBUG2);
                    break;
            }
        }
    }
};

/*!
 * Main application entry point. Creates an sbtool instance and lets it take over.
 */
int main(int argc, char *argv[], char *envp[])
{
    try
    {
        return encryptgpk(argc, argv).run();
    }
    catch (...)
    {
        Log::log(Logger::ERROR, "error: unexpected exception\n");
        return 1;
    }

    return 0;
}
