// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "connections/implementation/encryption_runner.h"

#include <cstdint>
#include <memory>
#include <string>
#include <utility>

#include "securegcm/ukey2_handshake.h"
#include "absl/strings/ascii.h"
#include "absl/time/time.h"
#include "connections/implementation/client_proxy.h"
#include "connections/implementation/endpoint_channel.h"
#include "internal/platform/base64_utils.h"
#include "internal/platform/byte_array.h"
#include "internal/platform/cancelable_alarm.h"
#include "internal/platform/exception.h"
#include "internal/platform/logging.h"
#include "internal/platform/scheduled_executor.h"

namespace nearby {
namespace connections {
namespace {

constexpr absl::Duration kTimeout = absl::Seconds(15);
constexpr std::int32_t kMaxUkey2VerificationStringLength = 32;
constexpr std::int32_t kTokenLength = 5;
constexpr securegcm::UKey2Handshake::HandshakeCipher kCipher =
    securegcm::UKey2Handshake::HandshakeCipher::P256_SHA512;

// Transforms a raw UKEY2 token (which is a random ByteArray that's
// kMaxUkey2VerificationStringLength long) into a kTokenLength string that only
// uses [A-Z], [0-9], '_', '-' for each character.
std::string ToHumanReadableString(const ByteArray& token) {
  std::string result = Base64Utils::Encode(token).substr(0, kTokenLength);
  absl::AsciiStrToUpper(&result);
  return result;
}

bool HandleEncryptionSuccess(const std::string& endpoint_id,
                             std::unique_ptr<securegcm::UKey2Handshake> ukey2,
                             EncryptionRunner::ResultListener& listener) {
  std::unique_ptr<std::string> verification_string =
      ukey2->GetVerificationString(kMaxUkey2VerificationStringLength);
  if (verification_string == nullptr) {
    return false;
  }

  ByteArray raw_authentication_token(*verification_string);

  listener.CallSuccessCallback(endpoint_id, std::move(ukey2),
                               ToHumanReadableString(raw_authentication_token),
                               raw_authentication_token);

  return true;
}

void CancelableAlarmRunnable(
    ClientProxy* client, const std::string& endpoint_id,
    std::shared_ptr<EndpointChannel> endpoint_channel) {
  LOG(INFO) << "Timing out encryption for client " << client->GetClientId()
            << " to endpoint_id=" << endpoint_id << " after "
            << absl::FormatDuration(kTimeout);
  endpoint_channel->Close();
}

class ServerRunnable final {
 public:
  ServerRunnable(ClientProxy* client, ScheduledExecutor* alarm_executor,
                 const std::string& endpoint_id,
                 std::shared_ptr<EndpointChannel> channel,
                 EncryptionRunner::ResultListener listener)
      : client_(client),
        alarm_executor_(alarm_executor),
        endpoint_id_(endpoint_id),
        weak_channel_(channel),
        listener_(std::move(listener)) {}

  void operator()() {
    // Lock the weak pointer. If it fails, the channel was freed.
    auto channel = weak_channel_.lock();
    // The IsClosed() check is to provide an early exit if channel has been
    // closed. Otherwise the Read() and Write() calls on the channel below will
    // return error and exit.
    if (!channel || channel->IsClosed()) {
      return;
    }
    CancelableAlarm timeout_alarm(
        "EncryptionRunner.StartServer() timeout",
        [this, weak_channel = weak_channel_]() {
          if (auto channel = weak_channel.lock()) {
            CancelableAlarmRunnable(client_, endpoint_id_, channel);
          }
        },
        kTimeout, alarm_executor_);

    std::unique_ptr<securegcm::UKey2Handshake> server =
        securegcm::UKey2Handshake::ForResponder(kCipher);
    if (server == nullptr) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    // Message 1 (Client Init)
    ExceptionOr<ByteArray> client_init = channel->Read();
    if (!client_init.ok()) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    securegcm::UKey2Handshake::ParseResult parse_result =
        server->ParseHandshakeMessage(std::string(client_init.result()));

    // Java code throws a HandshakeException / AlertException.
    if (!parse_result.success) {
      LogException();
      if (parse_result.alert_to_send != nullptr) {
        HandleAlertException(parse_result, channel);
      }
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    LOG(INFO) << "In StartServer(), read UKEY2 Message 1 from endpoint(id="
              << endpoint_id_ << ").";

    // Message 2 (Server Init)
    std::unique_ptr<std::string> server_init =
        server->GetNextHandshakeMessage();

    // Java code throws a HandshakeException.
    if (server_init == nullptr) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    Exception write_exception = channel->Write(*server_init);
    if (!write_exception.Ok()) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    LOG(INFO) << "In StartServer(), wrote UKEY2 Message 2 to endpoint(id="
              << endpoint_id_ << ").";

    // Message 3 (Client Finish)
    ExceptionOr<ByteArray> client_finish = channel->Read();

    if (!client_finish.ok()) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    parse_result =
        server->ParseHandshakeMessage(std::string(client_finish.result()));

    // Java code throws an AlertException or a HandshakeException.
    if (!parse_result.success) {
      LogException();
      if (parse_result.alert_to_send != nullptr) {
        HandleAlertException(parse_result, channel);
      }
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    LOG(INFO) << "In StartServer(), read UKEY2 Message 3 from endpoint(id="
              << endpoint_id_ << ").";

    timeout_alarm.Cancel();

    if (!HandleEncryptionSuccess(endpoint_id_, std::move(server), listener_)) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }
  }

 private:
  void LogException() const {
    LOG(ERROR) << "In StartServer(), UKEY2 failed with endpoint(id="
               << endpoint_id_ << ").";
  }

  void HandleHandshakeOrIoException(CancelableAlarm* timeout_alarm) {
    timeout_alarm->Cancel();
    listener_.CallFailureCallback(endpoint_id_);
  }

  void HandleAlertException(
      const securegcm::UKey2Handshake::ParseResult& parse_result,
      std::shared_ptr<EndpointChannel> channel) const {
    Exception write_exception = channel->Write(*parse_result.alert_to_send);
    if (!write_exception.Ok()) {
      LOG(WARNING) << "In StartServer(), client " << client_->GetClientId()
                   << " failed to pass the alert error message to endpoint(id="
                   << endpoint_id_ << ").";
    }
  }

  ClientProxy* client_;
  ScheduledExecutor* alarm_executor_;
  const std::string endpoint_id_;
  std::weak_ptr<EndpointChannel> weak_channel_;
  EncryptionRunner::ResultListener listener_;
};

class ClientRunnable final {
 public:
  ClientRunnable(ClientProxy* client, ScheduledExecutor* alarm_executor,
                 const std::string& endpoint_id,
                 std::shared_ptr<EndpointChannel> channel,
                 EncryptionRunner::ResultListener listener)
      : client_(client),
        alarm_executor_(alarm_executor),
        endpoint_id_(endpoint_id),
        weak_channel_(channel),
        listener_(std::move(listener)) {}

  void operator()() {
    // Lock the weak pointer. If it fails, the channel was freed.
    auto channel = weak_channel_.lock();
    // The IsClosed() check is to provide an early exit if channel has been
    // closed. Otherwise the Read() and Write() calls on the channel below will
    // return error and exit.
    if (!channel || channel->IsClosed()) {
      return;
    }

    CancelableAlarm timeout_alarm(
        "EncryptionRunner.StartClient() timeout",
        [this, weak_channel = weak_channel_]() {
          if (auto channel = weak_channel.lock()) {
            CancelableAlarmRunnable(client_, endpoint_id_, channel);
          }
        },
        kTimeout, alarm_executor_);

    std::unique_ptr<securegcm::UKey2Handshake> crypto =
        securegcm::UKey2Handshake::ForInitiator(kCipher);

    // Java code throws a HandshakeException.
    if (crypto == nullptr) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    // Message 1 (Client Init)
    std::unique_ptr<std::string> client_init =
        crypto->GetNextHandshakeMessage();

    // Java code throws a HandshakeException.
    if (client_init == nullptr) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    Exception write_init_exception = channel->Write(*client_init);
    if (!write_init_exception.Ok()) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    LOG(INFO) << "In StartClient(), wrote UKEY2 Message 1 to endpoint(id="
              << endpoint_id_ << ").";

    // Message 2 (Server Init)
    ExceptionOr<ByteArray> server_init = channel->Read();

    if (!server_init.ok()) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    securegcm::UKey2Handshake::ParseResult parse_result =
        crypto->ParseHandshakeMessage(std::string(server_init.result()));

    // Java code throws an AlertException or a HandshakeException.
    if (!parse_result.success) {
      LogException();
      if (parse_result.alert_to_send != nullptr) {
        HandleAlertException(parse_result, channel);
      }
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    LOG(INFO) << "In StartClient(), read UKEY2 Message 2 from endpoint(id="
              << endpoint_id_ << ").";

    // Message 3 (Client Finish)
    std::unique_ptr<std::string> client_finish =
        crypto->GetNextHandshakeMessage();

    // Java code throws a HandshakeException.
    if (client_finish == nullptr) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    Exception write_finish_exception = channel->Write(*client_finish);
    if (!write_finish_exception.Ok()) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }

    LOG(INFO) << "In StartClient(), wrote UKEY2 Message 3 to endpoint(id="
              << endpoint_id_ << ").";

    timeout_alarm.Cancel();

    if (!HandleEncryptionSuccess(endpoint_id_, std::move(crypto), listener_)) {
      LogException();
      HandleHandshakeOrIoException(&timeout_alarm);
      return;
    }
  }

 private:
  void LogException() const {
    LOG(ERROR) << "In StartClient(), UKEY2 failed with endpoint(id="
               << endpoint_id_ << ").";
  }

  void HandleHandshakeOrIoException(CancelableAlarm* timeout_alarm) {
    timeout_alarm->Cancel();
    listener_.CallFailureCallback(endpoint_id_);
  }

  void HandleAlertException(
      const securegcm::UKey2Handshake::ParseResult& parse_result,
      std::shared_ptr<EndpointChannel> channel) const {
    Exception write_exception = channel->Write(*parse_result.alert_to_send);
    if (!write_exception.Ok()) {
      LOG(WARNING) << "In StartClient(), client " << client_->GetClientId()
                   << " failed to pass the alert error message to endpoint(id="
                   << endpoint_id_ << ").";
    }
  }

  ClientProxy* client_;
  ScheduledExecutor* alarm_executor_;
  const std::string endpoint_id_;
  std::weak_ptr<EndpointChannel> weak_channel_;
  EncryptionRunner::ResultListener listener_;
};

}  // namespace

EncryptionRunner::~EncryptionRunner() { Shutdown(); }

void EncryptionRunner::StartServer(
    ClientProxy* client, const std::string& endpoint_id,
    std::shared_ptr<EndpointChannel> endpoint_channel,
    EncryptionRunner::ResultListener listener) {
  ServerRunnable runnable(client, &alarm_executor_, endpoint_id,
                          endpoint_channel, std::move(listener));
  server_executor_.Execute("encryption-server", std::move(runnable));
}

void EncryptionRunner::StartClient(
    ClientProxy* client, const std::string& endpoint_id,
    std::shared_ptr<EndpointChannel> endpoint_channel,
    EncryptionRunner::ResultListener listener) {
  ClientRunnable runnable(client, &alarm_executor_, endpoint_id,
                          endpoint_channel, std::move(listener));
  client_executor_.Execute("encryption-client", std::move(runnable));
}

void EncryptionRunner::Shutdown() {
  if (is_stopped_.Set(true)) {
    return;
  }

  // Stop all the ongoing Runnables (as gracefully as possible).
  client_executor_.Shutdown();
  server_executor_.Shutdown();
  alarm_executor_.Shutdown();
}

void EncryptionRunner::ResultListener::CallSuccessCallback(
    const std::string& endpoint_id,
    std::unique_ptr<securegcm::UKey2Handshake> ukey2,
    const std::string& auth_token, const ByteArray& raw_auth_token) {
  if (on_success_cb) {
    std::move(on_success_cb)(endpoint_id, std::move(ukey2), auth_token,
                             raw_auth_token);
  }
  Reset();
}

void EncryptionRunner::ResultListener::CallFailureCallback(
    const std::string& endpoint_id) {
  if (on_failure_cb) {
    std::move(on_failure_cb)(endpoint_id);
  }
  Reset();
}

void EncryptionRunner::ResultListener::Reset() {
  on_success_cb = nullptr;
  on_failure_cb = nullptr;
}

}  // namespace connections
}  // namespace nearby