main.cpp

#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <algorithm>
#include <cstring>
#include <termios.h>
#include <unistd.h>
#include <cstdint>
#include <sys/mman.h>

// Pour le mot de passe
typedef std::vector<char> SecureString;

extern "C" {
    // Utilitaires & Clés
    void generate_random_bytes(uint8_t* out, size_t len);
    int derive_keys(const uint8_t* pwd, size_t pwd_len, const uint8_t* salt, uint8_t* keys_out);
    
    // Compression
    void* create_compressor(int level);
    int compress_update(void* state, const uint8_t* data_in, size_t len_in, uint8_t* out, size_t* out_len);
    void compress_finalize(void* state, uint8_t* out, size_t* out_len);
    
    void* create_decompressor();
    int decompress_update(void* state, const uint8_t* data_in, size_t len_in);
    int decompress_pull(void* state, uint8_t* out, size_t out_capacity, size_t* out_len);
    void decompress_finalize(void* state, uint8_t* out, size_t* out_len);

    // Chiffrement
    void* create_encryptor(const uint8_t* key, const uint8_t* iv);
    void encrypt_update(void* state, const uint8_t* data_in, size_t data_in_len, uint8_t* out, size_t* out_len);
    size_t encrypt_finalize(void* state, uint8_t* out);
    
    void* create_decryptor(const uint8_t* key, const uint8_t* iv);
    void decrypt_update(void* state, const uint8_t* data_in, size_t data_in_len, uint8_t* out, size_t* out_len);
    size_t decrypt_finalize(void* state, uint8_t* out);

    // Authentification (HMAC)
    void* hmac_create(const uint8_t* key);
    void hmac_update(void* state, const uint8_t* data, size_t len);
    void hmac_finalize(void* state, uint8_t* out);
}

// Nettoyage sécurisé de la mémoire (empêche le compilateur d'ignorer)
void secure_zero(void* v, size_t n) {
    volatile unsigned char* p = (volatile unsigned char*)v;
    while (n--) *p++ = 0;
}

void lock_mem(void* addr, size_t size) {
    if (mlock(addr, size) != 0) {
        // avertir si mlock échoue (souvent besoin de sudo ou limite ulimit)
        perror("mlock failed");
    }
}

void unlock_mem(void* addr, size_t size) {
    munlock(addr, size);
}

SecureString get_secure_password() {
    termios oldt;
    tcgetattr(STDIN_FILENO, &oldt);
    termios newt = oldt;
    newt.c_lflag &= ~ECHO;
    tcsetattr(STDIN_FILENO, TCSANOW, &newt);

    SecureString password;
    password.reserve(128); 
    lock_mem(password.data(), password.capacity());

    std::cout << "Mot de passe : ";
    char ch;
    while (std::cin.get(ch) && ch != '\n') {
        password.push_back(ch);
    }

    tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
    std::cout << std::endl;
    return password;
}

void clean_pwd(SecureString& pwd) {
    secure_zero(pwd.data(), pwd.size());
    unlock_mem(pwd.data(), pwd.capacity());
    pwd.clear();
    pwd.shrink_to_fit();
}

void chiffrer_fichier(const std::string& in, const std::string& out, int compression_level) {
    SecureString pwd = get_secure_password();
    uint8_t salt[16], iv[16], keys[64];
    
    generate_random_bytes(salt, 16);
    generate_random_bytes(iv, 16);
    
    derive_keys((uint8_t*)pwd.data(), pwd.size(), salt, keys);
    clean_pwd(pwd); 

    std::ifstream f_in(in, std::ios::binary);
    std::ofstream f_out(out, std::ios::binary);

    void* hmac = hmac_create(keys + 32);
    // On authentifie le Salt et l'IV dès le début
    f_out.write((char*)salt, 16); hmac_update(hmac, salt, 16);
    f_out.write((char*)iv, 16);   hmac_update(hmac, iv, 16);

    void* comp = create_compressor(compression_level);
    void* enc = create_encryptor(keys, iv);

    // Buffers dimensionnés pour Zstd et le padding AES
    uint8_t b_in[4096], b_comp[8192], b_enc[9000];
    size_t sz_c, sz_e;

    // --- BOUCLE DE COMPRESSION + CHIFFREMENT ---
    while (f_in.read((char*)b_in, 4096) || f_in.gcount() > 0) {
        size_t read_bytes = f_in.gcount();
        
        // Zstd peut produire plus ou moins de données que l'entrée
        compress_update(comp, b_in, read_bytes, b_comp, &sz_c);
        
        if (sz_c > 0) {
            encrypt_update(enc, b_comp, sz_c, b_enc, &sz_e);
            if (sz_e > 0) { 
                f_out.write((char*)b_enc, sz_e); 
                hmac_update(hmac, b_enc, sz_e); 
            }
        }
    }

    // --- FINALISATION COMPRESSION ---
    // On force Zstd à vider ses derniers blocs
    compress_finalize(comp, b_comp, &sz_c);
    if (sz_c > 0) {
        encrypt_update(enc, b_comp, sz_c, b_enc, &sz_e);
        if (sz_e > 0) {
            f_out.write((char*)b_enc, sz_e); 
            hmac_update(hmac, b_enc, sz_e);
        }
    }

    // --- FINALISATION CHIFFREMENT ---
    // Ajout du padding PKCS7 final
    size_t final_e = encrypt_finalize(enc, b_enc);
    if (final_e > 0) {
        f_out.write((char*)b_enc, final_e); 
        hmac_update(hmac, b_enc, final_e);
    }

    // Génération de la signature
    uint8_t sig[32];
    hmac_finalize(hmac, sig);
    f_out.write((char*)sig, 32);
    
    // Nettoyage final
    secure_zero(keys, 64); 
    std::cout << "[+] Chiffrement + Compression terminés." << std::endl;
}

void dechiffrer_fichier(const std::string& in, const std::string& out) {
    std::ifstream f_in(in, std::ios::binary | std::ios::ate);
    std::streamsize file_sz = f_in.tellg();
    if (file_sz < 64) { 
        std::cerr << "Fichier trop court." << std::endl; 
        return; 
    }

    uint8_t hmac_sig[32];
    f_in.seekg(-32, std::ios::end); f_in.read((char*)hmac_sig, 32);

    f_in.seekg(0, std::ios::beg);
    uint8_t salt[16], iv[16], keys[64];
    f_in.read((char*)salt, 16); f_in.read((char*)iv, 16);

    SecureString pwd = get_secure_password();
    derive_keys((uint8_t*)pwd.data(), pwd.size(), salt, keys);
    clean_pwd(pwd);

    void* hmac = hmac_create(keys + 32);
    hmac_update(hmac, salt, 16); hmac_update(hmac, iv, 16);

    void* dec = create_decryptor(keys, iv);
    void* decomp = create_decompressor();
    std::ofstream f_out(out, std::ios::binary);

    uint8_t b_in[4096], b_dec[4112], b_final[16384];
    size_t sz_d, sz_raw;
    std::streamsize total_read = 32;

    // --- BOUCLE DE LECTURE PRINCIPALE ---
    while (total_read < file_sz - 32) {
        size_t to_read = std::min((size_t)4096, (size_t)(file_sz - 32 - total_read));
        f_in.read((char*)b_in, to_read);
        hmac_update(hmac, b_in, to_read);

        decrypt_update(dec, b_in, to_read, b_dec, &sz_d);
        if (sz_d > 0) {
            decompress_update(decomp, b_dec, sz_d);
            while (decompress_pull(decomp, b_final, 16384, &sz_raw) > 0) {
                if (sz_raw > 0) f_out.write((char*)b_final, sz_raw);
            }
        }
        total_read += to_read;
    }

    // --- FINALISATION DÉCHIFFREMENT ---
    size_t last_d = decrypt_finalize(dec, b_dec);
    if (last_d > 0) {
        decompress_update(decomp, b_dec, last_d);
        // On vide tout ce qui sort suite au dernier bloc AES
        while (decompress_pull(decomp, b_final, 16384, &sz_raw) > 0) {
            if (sz_raw > 0) f_out.write((char*)b_final, sz_raw);
        }
    }
    
    // --- FINALISATION DÉCOMPRESSION ---
    // On fait un dernier pull au cas où Zstd aurait gardé des miettes
    while (decompress_pull(decomp, b_final, 16384, &sz_raw) > 0) {
        if (sz_raw > 0) f_out.write((char*)b_final, sz_raw);
    }

    // dernier appel, libère la mémoire côté rust
    decompress_finalize(decomp, b_final, &sz_raw);
    if (sz_raw > 0) f_out.write((char*)b_final, sz_raw);

    // Vérification hmac
    uint8_t hmac_calc[32];
    hmac_finalize(hmac, hmac_calc);

    bool ok = (std::memcmp(hmac_sig, hmac_calc, 32) == 0);
    secure_zero(keys, 64); // On efface les clés de la RAM

    if (!ok) {
        std::cerr << "[!] ALERTE : HMAC invalide !" << std::endl;
        f_out.close(); 
        std::remove(out.c_str()); // On supprime le fichier corrompu
    } else {
        std::cout << "[+] Décompression + Déchiffrement réussis." << std::endl;
    }
}

int main(int argc, char* argv[]) {
    if (argc < 4) { 
        std::cout << "Usage: " << argv[0] << " <e|d> <input_file> <output_file>" << std::endl; 
        return 1; 
    }
    
    std::string mode = argv[1];
    int level = 3;
    if (argc == 5 && mode == "e") {
        level = std::stoi(argv[4]);
    }
    if (mode == "e") {
        chiffrer_fichier(argv[2], argv[3], level);
    } else if (mode == "d") {
        dechiffrer_fichier(argv[2], argv[3]);
    } else {
        std::cerr << "Mode inconnu. Utilisez 'e' pour chiffrer + compression ou 'd' pour décompresser + déchiffrer." << std::endl;
        return 1;
    }
    
    return 0;
}