crossdesk/src/thumbnail/thumbnail.cpp

#include "thumbnail.h"

#include <openssl/aes.h>
#include <openssl/crypto.h>
#include <openssl/evp.h>
#include <openssl/rand.h>

#include <chrono>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <string>
#include <vector>

#include "libyuv.h"
#include "rd_log.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

namespace crossdesk {

void ScaleNv12ToABGR(char* src, int src_w, int src_h, int dst_w, int dst_h,
                     char* dst_rgba) {
  uint8_t* y = reinterpret_cast<uint8_t*>(src);
  uint8_t* uv = y + src_w * src_h;

  float src_aspect = float(src_w) / src_h;
  float dst_aspect = float(dst_w) / dst_h;
  int fit_w = dst_w, fit_h = dst_h;
  if (src_aspect > dst_aspect) {
    fit_h = int(dst_w / src_aspect);
  } else {
    fit_w = int(dst_h * src_aspect);
  }

  std::vector<uint8_t> y_i420(src_w * src_h);
  std::vector<uint8_t> u_i420((src_w / 2) * (src_h / 2));
  std::vector<uint8_t> v_i420((src_w / 2) * (src_h / 2));
  libyuv::NV12ToI420(y, src_w, uv, src_w, y_i420.data(), src_w, u_i420.data(),
                     src_w / 2, v_i420.data(), src_w / 2, src_w, src_h);

  std::vector<uint8_t> y_fit(fit_w * fit_h);
  std::vector<uint8_t> u_fit((fit_w + 1) / 2 * (fit_h + 1) / 2);
  std::vector<uint8_t> v_fit((fit_w + 1) / 2 * (fit_h + 1) / 2);
  libyuv::I420Scale(y_i420.data(), src_w, u_i420.data(), src_w / 2,
                    v_i420.data(), src_w / 2, src_w, src_h, y_fit.data(), fit_w,
                    u_fit.data(), (fit_w + 1) / 2, v_fit.data(),
                    (fit_w + 1) / 2, fit_w, fit_h, libyuv::kFilterBilinear);

  std::vector<uint8_t> abgr(fit_w * fit_h * 4);
  libyuv::I420ToABGR(y_fit.data(), fit_w, u_fit.data(), (fit_w + 1) / 2,
                     v_fit.data(), (fit_w + 1) / 2, abgr.data(), fit_w * 4,
                     fit_w, fit_h);

  std::memset(dst_rgba, 0, dst_w * dst_h * 4);
  for (int i = 0; i < dst_w * dst_h; ++i) {
    dst_rgba[i * 4 + 3] = static_cast<char>(0xFF);
  }

  for (int row = 0; row < fit_h; ++row) {
    int dst_offset =
        ((row + (dst_h - fit_h) / 2) * dst_w + (dst_w - fit_w) / 2) * 4;
    std::memcpy(dst_rgba + dst_offset, abgr.data() + row * fit_w * 4,
                fit_w * 4);
  }
}

Thumbnail::Thumbnail(std::string save_path) {
  if (!save_path.empty()) {
    save_path_ = save_path;
  }

  RAND_bytes(aes128_key_, sizeof(aes128_key_));
  RAND_bytes(aes128_iv_, sizeof(aes128_iv_));
  std::filesystem::create_directories(save_path_);
}

Thumbnail::Thumbnail(std::string save_path, unsigned char* aes128_key,
                     unsigned char* aes128_iv) {
  if (!save_path.empty()) {
    save_path_ = save_path;
  }

  std::memcpy(aes128_key_, aes128_key, sizeof(aes128_key_));
  std::memcpy(aes128_iv_, aes128_iv, sizeof(aes128_iv_));
  std::filesystem::create_directories(save_path_);
}

Thumbnail::~Thumbnail() {
  if (rgba_buffer_) {
    delete[] rgba_buffer_;
    rgba_buffer_ = nullptr;
  }
}

int Thumbnail::SetThumbnailDpiScale(float dpi_scale) {
  thumbnail_width_ = static_cast<int>(thumbnail_width_ * dpi_scale);
  thumbnail_height_ = static_cast<int>(thumbnail_height_ * dpi_scale);
  return 0;
}

int Thumbnail::SaveToThumbnail(const char* yuv420p, int width, int height,
                               const std::string& remote_id,
                               const std::string& host_name,
                               const std::string& password) {
  if (!rgba_buffer_) {
    rgba_buffer_ = new char[thumbnail_width_ * thumbnail_height_ * 4];
  }

  if (yuv420p) {
    ScaleNv12ToABGR((char*)yuv420p, width, height, thumbnail_width_,
                    thumbnail_height_, rgba_buffer_);
  } else {
    // If yuv420p is null, fill the buffer with black pixels
    std::memset(rgba_buffer_, 0x00, thumbnail_width_ * thumbnail_height_ * 4);
    for (int i = 0; i < thumbnail_width_ * thumbnail_height_; ++i) {
      // Set alpha channel to opaque
      rgba_buffer_[i * 4 + 3] = static_cast<char>(0xFF);
    }
  }

  std::string image_file_name;
  if (password.empty()) {
    return 0;
  } else {
    // delete the old thumbnail
    std::string filename_with_remote_id = remote_id;
    DeleteThumbnail(filename_with_remote_id);
  }

  std::string cipher_password = AES_encrypt(password, aes128_key_, aes128_iv_);
  image_file_name = remote_id + 'Y' + host_name + '@' + cipher_password;
  std::string file_path = save_path_ + image_file_name;
  stbi_write_png(file_path.data(), thumbnail_width_, thumbnail_height_, 4,
                 rgba_buffer_, thumbnail_width_ * 4);

  return 0;
}

int Thumbnail::LoadThumbnail(
    std::vector<std::pair<std::string, Thumbnail::RecentConnection>>&
        recent_connections,
    int* width, int* height) {
  recent_connections.clear();
  if (width) *width = thumbnail_width_;
  if (height) *height = thumbnail_height_;

  std::vector<std::filesystem::path> image_paths =
      FindThumbnailPath(save_path_);

  if (image_paths.size() == 0) {
    return -1;
  } else {
    for (int i = 0; i < image_paths.size(); i++) {
      // size_t pos1 = image_paths[i].string().find('/') + 1;
      std::string cipher_image_name = image_paths[i].filename().string();
      std::string remote_id;
      std::string cipher_password;
      std::string remote_host_name;
      std::string password;
      std::string original_image_name;
      bool remember_password = false;

      if (cipher_image_name.size() > 9 && 'Y' == cipher_image_name[9] &&
          cipher_image_name.size() >= 16) {
        size_t pos_y = cipher_image_name.find('Y');
        size_t pos_at = cipher_image_name.find('@');

        if (pos_y == std::string::npos || pos_at == std::string::npos ||
            pos_y >= pos_at) {
          LOG_ERROR("Invalid filename");
          continue;
        }

        remote_id = cipher_image_name.substr(0, pos_y);
        remote_host_name =
            cipher_image_name.substr(pos_y + 1, pos_at - pos_y - 1);
        cipher_password = cipher_image_name.substr(pos_at + 1);
        password = AES_decrypt(cipher_password, aes128_key_, aes128_iv_);
        remember_password = true;

        original_image_name = remote_id + 'Y' + remote_host_name + "@" +
                              password;
      } else {
        size_t pos_n = cipher_image_name.find('N');
        // size_t pos_at = cipher_image_name.find('@');

        if (pos_n == std::string::npos) {
          LOG_ERROR("Invalid filename");
          continue;
        }

        remote_id = cipher_image_name.substr(0, pos_n);
        remote_host_name = cipher_image_name.substr(pos_n + 1);

        original_image_name = remote_id + 'N' + remote_host_name;
      }

      std::string image_path = save_path_ + cipher_image_name;
      Thumbnail::RecentConnection recent_connection;
      recent_connection.remote_id = remote_id;
      recent_connection.remote_host_name = remote_host_name;
      recent_connection.password = password;
      recent_connection.remember_password = remember_password;
      recent_connection.image_path = image_path;
      recent_connections.emplace_back(
          std::make_pair(original_image_name, recent_connection));
    }
    return 0;
  }
  return 0;
}

bool Thumbnail::DecodeImage(const std::filesystem::path& image_path,
                            std::vector<unsigned char>* rgba, int* width,
                            int* height) const {
  if (!rgba || !width || !height) {
    return false;
  }

  int channels = 0;
  unsigned char* pixels =
      stbi_load(image_path.string().c_str(), width, height, &channels, 4);
  if (!pixels || *width <= 0 || *height <= 0) {
    if (pixels) {
      stbi_image_free(pixels);
    }
    rgba->clear();
    return false;
  }

  const size_t byte_count =
      static_cast<size_t>(*width) * static_cast<size_t>(*height) * 4;
  rgba->assign(pixels, pixels + byte_count);
  stbi_image_free(pixels);
  return true;
}

int Thumbnail::DeleteThumbnail(const std::string& filename_keyword) {
  for (const auto& entry : std::filesystem::directory_iterator(save_path_)) {
    if (entry.is_regular_file()) {
      const std::string filename = entry.path().filename().string();
      std::string id_hostname = filename_keyword.substr(0, filename.find('@'));
      if (filename.find(id_hostname) != std::string::npos) {
        std::filesystem::remove(entry.path());
      }
    }
  }

  return 0;
}

std::vector<std::filesystem::path> Thumbnail::FindThumbnailPath(
    const std::filesystem::path& directory) {
  std::vector<std::filesystem::path> thumbnails_path;

  if (!std::filesystem::is_directory(directory)) {
    LOG_ERROR("No such directory [{}]", directory.string());
    return thumbnails_path;
  }

  for (const auto& entry : std::filesystem::directory_iterator(directory)) {
    if (entry.is_regular_file()) {
      thumbnails_path.push_back(entry.path());
    }
  }

  std::sort(thumbnails_path.begin(), thumbnails_path.end(),
            [](const std::filesystem::path& a, const std::filesystem::path& b) {
              return std::filesystem::last_write_time(a) >
                     std::filesystem::last_write_time(b);
            });

  return thumbnails_path;
}

int Thumbnail::DeleteAllFilesInDirectory() {
  if (std::filesystem::exists(save_path_) &&
      std::filesystem::is_directory(save_path_)) {
    for (const auto& entry : std::filesystem::directory_iterator(save_path_)) {
      if (std::filesystem::is_regular_file(entry.status())) {
        std::filesystem::remove(entry.path());
      }
    }
    return 0;
  }
  return -1;
}

std::string Thumbnail::AES_encrypt(const std::string& plaintext,
                                   unsigned char* key, unsigned char* iv) {
  EVP_CIPHER_CTX* ctx;
  int len;
  int ciphertext_len;
  int ret = 0;
  std::vector<unsigned char> ciphertext(plaintext.size() + AES_BLOCK_SIZE);

  ctx = EVP_CIPHER_CTX_new();
  if (!ctx) {
    LOG_ERROR("Error in EVP_CIPHER_CTX_new");
    return plaintext;
  }

  ret = EVP_EncryptInit_ex(ctx, EVP_aes_128_cbc(), NULL, key, iv);
  if (1 != ret) {
    LOG_ERROR("Error in EVP_EncryptInit_ex");
    EVP_CIPHER_CTX_free(ctx);
    return plaintext;
  }

  ret = EVP_EncryptUpdate(
      ctx, ciphertext.data(), &len,
      reinterpret_cast<const unsigned char*>(plaintext.data()),
      (int)plaintext.size());
  if (1 != ret) {
    LOG_ERROR("Error in EVP_EncryptUpdate");
    EVP_CIPHER_CTX_free(ctx);
    return plaintext;
  }

  ciphertext_len = len;
  ret = EVP_EncryptFinal_ex(ctx, ciphertext.data() + len, &len);
  if (1 != ret) {
    LOG_ERROR("Error in EVP_EncryptFinal_ex");
    EVP_CIPHER_CTX_free(ctx);
    return plaintext;
  }

  ciphertext_len += len;

  unsigned char hex_str[256];
  size_t hex_str_len = 0;
  ret = OPENSSL_buf2hexstr_ex((char*)hex_str, sizeof(hex_str), &hex_str_len,
                              ciphertext.data(), ciphertext_len, '\0');
  if (1 != ret) {
    LOG_ERROR("Error in OPENSSL_buf2hexstr_ex");
    EVP_CIPHER_CTX_free(ctx);
    return plaintext;
  }

  EVP_CIPHER_CTX_free(ctx);

  std::string str(reinterpret_cast<char*>(hex_str), hex_str_len);
  return str;
}

std::string Thumbnail::AES_decrypt(const std::string& ciphertext,
                                   unsigned char* key, unsigned char* iv) {
  unsigned char ciphertext_buf[256];
  size_t ciphertext_buf_len = 0;
  unsigned char plaintext[256];
  int plaintext_len = 0;
  int plaintext_final_len = 0;
  EVP_CIPHER_CTX* ctx;
  int ret = 0;

  ret = OPENSSL_hexstr2buf_ex(ciphertext_buf, sizeof(ciphertext_buf),
                              &ciphertext_buf_len, ciphertext.c_str(), '\0');
  if (1 != ret) {
    LOG_ERROR("Error in OPENSSL_hexstr2buf_ex");
    return ciphertext;
  }

  ctx = EVP_CIPHER_CTX_new();
  if (!ctx) {
    LOG_ERROR("Error in EVP_CIPHER_CTX_new");
    return ciphertext;
  }

  ret = EVP_DecryptInit_ex(ctx, EVP_aes_128_cbc(), NULL, key, iv);
  if (1 != ret) {
    LOG_ERROR("Error in EVP_DecryptInit_ex");

    EVP_CIPHER_CTX_free(ctx);
    return ciphertext;
  }

  ret = EVP_DecryptUpdate(ctx, plaintext, &plaintext_len, ciphertext_buf,
                          (int)ciphertext_buf_len);
  if (1 != ret) {
    LOG_ERROR("Error in EVP_DecryptUpdate");

    EVP_CIPHER_CTX_free(ctx);
    return ciphertext;
  }

  ret =
      EVP_DecryptFinal_ex(ctx, plaintext + plaintext_len, &plaintext_final_len);
  if (1 != ret) {
    LOG_ERROR("Error in EVP_DecryptFinal_ex");

    EVP_CIPHER_CTX_free(ctx);
    return ciphertext;
  }
  plaintext_len += plaintext_final_len;

  EVP_CIPHER_CTX_free(ctx);

  return std::string(reinterpret_cast<char*>(plaintext), plaintext_len);
}
}  // namespace crossdesk