// Copyright 2013 Google Inc. All Rights Reserved.

#include "device_files.h"

#include <cstring>
#include <string>
#include <sstream>
#include <unistd.h>

#include "device_files.pb.h"
#include "file_store.h"
#include "log.h"
#include "openssl/sha.h"

namespace wvcdm {

// TODO(rfrias): Make this work for non-unix paths
const char* DeviceFiles::kBasePath = "/data/mediadrm/IDM";
const char* DeviceFiles::kPathDelimiter = "/";
const char* DeviceFiles::kDeviceCertificateFileName = "cert.bin";

// Protobuf generated classes.
using video_widevine_client::sdk::DeviceCertificate;
using video_widevine_client::sdk::HashedFile;

bool DeviceFiles::StoreCertificate(const std::string& certificate,
                                   const std::string& wrapped_private_key) {
  // Fill in file information
  video_widevine_client::sdk::File file;

  file.set_type(video_widevine_client::sdk::File::DEVICE_CERTIFICATE);
  file.set_version(video_widevine_client::sdk::File::VERSION_1);

  DeviceCertificate *device_certificate = file.mutable_device_certificate();
  device_certificate->set_certificate(certificate);
  device_certificate->set_wrapped_private_key(wrapped_private_key);

  std::string serialized_string;
  file.SerializeToString(&serialized_string);

  // calculate SHA hash
  std::string hash;
  if (!Hash(serialized_string, &hash)) {
    LOGW("DeviceFiles::StoreCertificate: Hash computation failed");
    return false;
  }

  // File in hashed file data
  HashedFile hashed_file;
  hashed_file.set_file(serialized_string);
  hashed_file.set_hash(hash);

  hashed_file.SerializeToString(&serialized_string);

  return StoreFile(kDeviceCertificateFileName, serialized_string);
}

bool DeviceFiles::RetrieveCertificate(std::string* certificate,
                                      std::string* wrapped_private_key) {

  std::string serialized_hashed_file;
  if (!RetrieveFile(kDeviceCertificateFileName, &serialized_hashed_file))
    return false;

  HashedFile hashed_file;
  if (!hashed_file.ParseFromString(serialized_hashed_file)) {
    LOGW("DeviceFiles::RetrieveCertificate: Unable to parse hash file");
    return false;
  }

  std::string hash;
  if (!Hash(hashed_file.file(), &hash)) {
    LOGW("DeviceFiles::RetrieveCertificate: Hash computation failed");
    return false;
  }

  if (hash.compare(hashed_file.hash())) {
    LOGW("DeviceFiles::RetrieveCertificate: Hash mismatch");
    return false;
  }

  video_widevine_client::sdk::File file;
  if (!file.ParseFromString(hashed_file.file())) {
    LOGW("DeviceFiles::RetrieveCertificate: Unable to parse file");
    return false;
  }

  if (file.type() != video_widevine_client::sdk::File::DEVICE_CERTIFICATE) {
    LOGW("DeviceFiles::RetrieveCertificate: Incorrect file type");
    return false;
  }

  if (file.version() != video_widevine_client::sdk::File::VERSION_1) {
    LOGW("DeviceFiles::RetrieveCertificate: Incorrect file version");
    return false;
  }

  if (!file.has_device_certificate()) {
    LOGW("DeviceFiles::RetrieveCertificate: Certificate not present");
    return false;
  }

  DeviceCertificate device_certificate = file.device_certificate();

  *certificate = device_certificate.certificate();
  *wrapped_private_key = device_certificate.wrapped_private_key();
  return true;
}

bool DeviceFiles::Hash(const std::string& data, std::string* hash) {
  if (!hash)
    return false;

  hash->resize(SHA256_DIGEST_LENGTH);
  SHA256_CTX sha256;
  SHA256_Init(&sha256);
  SHA256_Update(&sha256, data.data(), data.size());
  SHA256_Final(reinterpret_cast<unsigned char*>(const_cast<char*>(hash->data())), &sha256);
  return true;
}

bool DeviceFiles::StoreFile(const char* name, const std::string& data) {
  if (!name)
    return false;

  std::string path = GetBasePath(kBasePath);

  if (!File::IsDirectory(path)) {
    if (!File::CreateDirectory(path))
      return false;
  }

  path += name;

  File file(path, File::kCreate | File::kTruncate | File::kBinary);
  if (file.IsBad()) {
    LOGW("DeviceFiles::StoreFile: File open failed: %s", path.c_str());
    return false;
  }

  ssize_t bytes = file.Write(data.data(), data.size());

  file.Close();

  if (bytes != static_cast<ssize_t>(data.size())) {
    LOGW("DeviceFiles::StoreFile: write failed: %d %d", data.size(), bytes);
    return false;
  }

  return true;
}

bool DeviceFiles::RetrieveFile(const char* name, std::string* data) {
  if (!data)
    return false;

  std::string path = GetBasePath(kBasePath) + name;

  if (!File::Exists(path)) {
    LOGW("DeviceFiles::RetrieveFile: %s does not exist", path.c_str());
    return false;
  }

  ssize_t bytes = File::FileSize(path);

  if (bytes <= 0) {
    LOGW("DeviceFiles::RetrieveFile: File size invalid: %d", path.c_str());
    return false;
  }

  File file(path, File::kReadOnly | File::kBinary);
  if (file.IsBad()) {
    LOGW("DeviceFiles::RetrieveFile: File open failed: %s", path.c_str());
    return false;
  }

  data->resize(bytes);

  bytes = file.Read(reinterpret_cast<void*>(const_cast<char*>(data->data())),
      data->size());

  if (bytes != static_cast<ssize_t>(data->size())) {
    LOGW("DeviceFiles::StoreFile: write failed: %d %d", data->size(), bytes);
    return false;
  }

  return true;
}

std::string DeviceFiles::GetBasePath(const char* dir) {
  // TODO(rfrias): Make this work for non-unix paths
  std::stringstream ss;
  ss << dir << getuid() << kPathDelimiter;
  return ss.str();
}

}