// Copyright 2022 The Chromium Authors.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <time.h>

#include <iostream>
#include <memory>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <vector>

#include "content_analysis/sdk/analysis_client.h"
#include "demo/atomic_output.h"

using content_analysis::sdk::Client;
using content_analysis::sdk::ContentAnalysisRequest;
using content_analysis::sdk::ContentAnalysisResponse;
using content_analysis::sdk::ContentAnalysisAcknowledgement;

// Different paths are used depending on whether this agent should run as a
// use specific agent or not.  These values are chosen to match the test
// values in chrome browser.
constexpr char kPathUser[] = "path_user";
constexpr char kPathSystem[] = "brcm_chrm_cas";

// Global app config.
std::string path = kPathSystem;
bool user_specific = false;
bool group = false;
std::unique_ptr<Client> client;

// Paramters used to build the request.
content_analysis::sdk::AnalysisConnector connector =
    content_analysis::sdk::FILE_ATTACHED;
time_t request_token_number = time(nullptr);
std::string request_token;
std::string tag = "dlp";
bool threaded = false;
std::string digest = "sha256-123456";
std::string url = "https://upload.example.com";
std::string email = "me@example.com";
std::string machine_user = "DOMAIN\\me";
std::vector<std::string> datas;

// When grouping, remember the tokens of all requests/responses in order to
// acknowledge them all with the same final action.
//
// This global state.  It may be access from multiple thread so must be
// accessed from a critical section.
std::mutex global_mutex;
ContentAnalysisAcknowledgement::FinalAction global_final_action =
    ContentAnalysisAcknowledgement::ALLOW;
std::vector<std::string> request_tokens;

// Command line parameters.
constexpr const char* kArgConnector = "--connector=";
constexpr const char* kArgDigest = "--digest=";
constexpr const char* kArgEmail = "--email=";
constexpr const char* kArgGroup = "--group";
constexpr const char* kArgMachineUser = "--machine-user=";
constexpr const char* kArgPath = "--path=";
constexpr const char* kArgRequestToken = "--request-token=";
constexpr const char* kArgTag = "--tag=";
constexpr const char* kArgThreaded = "--threaded";
constexpr const char* kArgUrl = "--url=";
constexpr const char* kArgUserSpecific = "--user";
constexpr const char* kArgHelp = "--help";

bool ParseCommandLine(int argc, char* argv[]) {
  for (int i = 1; i < argc; ++i) {
    const std::string arg = argv[i];
    if (arg.find(kArgConnector) == 0) {
      std::string connector_str = arg.substr(strlen(kArgConnector));
      if (connector_str == "download") {
        connector = content_analysis::sdk::FILE_DOWNLOADED;
      } else if (connector_str == "attach") {
        connector = content_analysis::sdk::FILE_ATTACHED;
      } else if (connector_str == "bulk-data-entry") {
        connector = content_analysis::sdk::BULK_DATA_ENTRY;
      } else if (connector_str == "print") {
        connector = content_analysis::sdk::PRINT;
      } else if (connector_str == "file-transfer") {
        connector = content_analysis::sdk::FILE_TRANSFER;
      } else if (connector_str == "data-copied") {
        connector = content_analysis::sdk::DATA_COPIED;
      } else {
        std::cout << "[Demo] Incorrect command line arg: " << arg << std::endl;
        return false;
      }
    } else if (arg.find(kArgRequestToken) == 0) {
      request_token = arg.substr(strlen(kArgRequestToken));
    } else if (arg.find(kArgTag) == 0) {
      tag = arg.substr(strlen(kArgTag));
    } else if (arg.find(kArgThreaded) == 0) {
      threaded = true;
    } else if (arg.find(kArgDigest) == 0) {
      digest = arg.substr(strlen(kArgDigest));
    } else if (arg.find(kArgUrl) == 0) {
      url = arg.substr(strlen(kArgUrl));
    } else if (arg.find(kArgMachineUser) == 0) {
      machine_user = arg.substr(strlen(kArgMachineUser));
    } else if (arg.find(kArgEmail) == 0) {
      email = arg.substr(strlen(kArgEmail));
    } else if (arg.find(kArgPath) == 0) {
      path = arg.substr(strlen(kArgPath));
    } else if (arg.find(kArgUserSpecific) == 0) {
      // If kArgPath was already used, abort.
      if (path != kPathSystem) {
        std::cout << std::endl << "ERROR: use --path=<path> after --user";
        return false;
      }
      path = kPathUser;
      user_specific = true;
    } else if (arg.find(kArgGroup) == 0) {
      group = true;
    } else if (arg.find(kArgHelp) == 0) {
      return false;
    } else {
      datas.push_back(arg);
    }
  }

  return true;
}

void PrintHelp() {
  std::cout
    << std::endl << std::endl
    << "Usage: client [OPTIONS] [@]content_or_file ..." << std::endl
    << "A simple client to send content analysis requests to a running agent." << std::endl
    << "Without @ the content to analyze is the argument itself." << std::endl
    << "Otherwise the content is read from a file called 'content_or_file'." << std::endl
    << "Multiple [@]content_or_file arguments may be specified, each generates one request." << std::endl
    << std::endl << "Options:"  << std::endl
    << kArgConnector << "<connector> : one of 'download', 'attach' (default), 'bulk-data-entry', 'print', 'file-transfer', or 'data-copied'" << std::endl
    << kArgRequestToken << "<unique-token> : defaults to 'req-<number>' which auto increments" << std::endl
    << kArgTag << "<tag> : defaults to 'dlp'" << std::endl
    << kArgThreaded << " : handled multiple requests using threads" << std::endl
    << kArgUrl << "<url> : defaults to 'https://upload.example.com'" << std::endl
    << kArgMachineUser << "<machine-user> : defaults to 'DOMAIN\\me'" << std::endl
    << kArgEmail << "<email> : defaults to 'me@example.com'" << std::endl
    << kArgPath << " <path> : Used the specified path instead of default. Must come after --user." << std::endl
    << kArgUserSpecific << " : Connects to an OS user specific agent" << std::endl
    << kArgDigest << "<digest> : defaults to 'sha256-123456'" << std::endl
    << kArgGroup << " : Generate the same final action for all requests" << std::endl
    << kArgHelp << " : prints this help message" << std::endl;
}

std::string GenerateRequestToken() {
  std::stringstream stm;
  stm << "req-" << request_token_number++;
  return stm.str();
}

ContentAnalysisRequest BuildRequest(const std::string& data) {
  std::string filepath;
  std::string filename;
  if (data[0] == '@') {
    filepath = data.substr(1);
    filename = filepath.substr(filepath.find_last_of("/\\") + 1);
  }

  ContentAnalysisRequest request;

  // Set request to expire 5 minutes into the future.
  request.set_expires_at(time(nullptr) + 5 * 60);
  request.set_analysis_connector(connector);
  request.set_request_token(!request_token.empty()
      ? request_token : GenerateRequestToken());
  *request.add_tags() = tag;

  auto request_data = request.mutable_request_data();
  request_data->set_url(url);
  request_data->set_email(email);
  request_data->set_digest(digest);
  if (!filename.empty()) {
    request_data->set_filename(filename);
  }

  auto client_metadata = request.mutable_client_metadata();
  auto browser = client_metadata->mutable_browser();
  browser->set_machine_user(machine_user);

  if (!filepath.empty()) {
    request.set_file_path(filepath);
  } else if (!data.empty()) {
    request.set_text_content(data);
  } else {
    std::cout << "[Demo] Specify text content or a file path." << std::endl;
    PrintHelp();
    exit(1);
  }

  return request;
}

// Gets the most severe action within the result.
ContentAnalysisResponse::Result::TriggeredRule::Action
GetActionFromResult(const ContentAnalysisResponse::Result& result) {
  auto action =
    ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED;
  for (auto rule : result.triggered_rules()) {
    if (rule.has_action() && rule.action() > action)
      action = rule.action();
  }
  return action;
}

// Gets the most severe action within all the the results of a response.
ContentAnalysisResponse::Result::TriggeredRule::Action
GetActionFromResponse(const ContentAnalysisResponse& response) {
  auto action =
    ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED;
  for (auto result : response.results()) {
    auto action2 = GetActionFromResult(result);
    if (action2 > action)
      action = action2;
  }
  return action;
}

void DumpResponse(
    std::stringstream& stream,
    const ContentAnalysisResponse& response) {
  for (auto result : response.results()) {
    auto tag = result.has_tag() ? result.tag() : "<no-tag>";

    auto status = result.has_status()
      ? result.status()
      : ContentAnalysisResponse::Result::STATUS_UNKNOWN;
    std::string status_str;
    switch (status) {
    case ContentAnalysisResponse::Result::STATUS_UNKNOWN:
      status_str = "Unknown";
      break;
    case ContentAnalysisResponse::Result::SUCCESS:
      status_str = "Success";
      break;
    case ContentAnalysisResponse::Result::FAILURE:
      status_str = "Failure";
      break;
    default:
      status_str = "<Uknown>";
      break;
    }

    auto action = GetActionFromResult(result);
    std::string action_str;
    switch (action) {
    case ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED:
      action_str = "allowed";
      break;
    case ContentAnalysisResponse::Result::TriggeredRule::REPORT_ONLY:
      action_str = "reported only";
      break;
    case ContentAnalysisResponse::Result::TriggeredRule::WARN:
      action_str = "warned";
      break;
    case ContentAnalysisResponse::Result::TriggeredRule::BLOCK:
      action_str = "blocked";
      break;
    }

    time_t now = time(nullptr);
    stream << "[Demo] Request " << response.request_token()  << " is " << action_str
           << " after " << tag
           << " analysis, status=" << status_str
           << " at " << ctime(&now);
  }
}

ContentAnalysisAcknowledgement BuildAcknowledgement(
    const std::string& request_token,
    ContentAnalysisAcknowledgement::FinalAction final_action) {
  ContentAnalysisAcknowledgement ack;
  ack.set_request_token(request_token);
  ack.set_status(ContentAnalysisAcknowledgement::SUCCESS);
  ack.set_final_action(final_action);
  return ack;
}

void HandleRequest(const ContentAnalysisRequest& request) {
  AtomicCout aout;
  ContentAnalysisResponse response;
  int err = client->Send(request, &response);
  if (err != 0) {
    aout.stream() << "[Demo] Error sending request " << request.request_token()
           << std::endl;
  } else if (response.results_size() == 0) {
    aout.stream() << "[Demo] Response " << request.request_token() << " is missing a result"
                  << std::endl;
  } else {
    DumpResponse(aout.stream(), response);

    auto final_action = ContentAnalysisAcknowledgement::ALLOW;
    switch (GetActionFromResponse(response)) {
    case ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED:
      break;
    case ContentAnalysisResponse::Result::TriggeredRule::REPORT_ONLY:
      final_action = ContentAnalysisAcknowledgement::REPORT_ONLY;
      break;
    case ContentAnalysisResponse::Result::TriggeredRule::WARN:
      final_action = ContentAnalysisAcknowledgement::WARN;
      break;
    case ContentAnalysisResponse::Result::TriggeredRule::BLOCK:
      final_action = ContentAnalysisAcknowledgement::BLOCK;
      break;
    }

    // If grouping, remember the request's token in order to ack the response
    // later.
    if (group) {
      std::unique_lock<std::mutex> lock(global_mutex);
      request_tokens.push_back(request.request_token());
      if (final_action > global_final_action)
        global_final_action = final_action;
    } else {
      int err = client->Acknowledge(
        BuildAcknowledgement(response.request_token(), final_action));
      if (err != 0) {
        aout.stream() << "[Demo] Error sending ack " << request.request_token()
                      << std::endl;
      }
    }
  }
}

void ProcessRequest(size_t i) {
  auto request = BuildRequest(datas[i]);

  {
    AtomicCout aout;
    aout.stream() << "[Demo] Sending request " << request.request_token() << std::endl;
  }

  HandleRequest(request);
}

int main(int argc, char* argv[]) {
  if (!ParseCommandLine(argc, argv)) {
    PrintHelp();
    return 1;
  }

  // Each client uses a unique name to identify itself with Google Chrome.
  client = Client::Create({path, user_specific});
  if (!client) {
    std::cout << "[Demo] Error starting client" << std::endl;
    return 1;
  };

  auto info = client->GetAgentInfo();
  std::cout << "Agent pid=" << info.pid
            << " path=" << info.binary_path << std::endl;

  if (threaded) {
    std::vector<std::unique_ptr<std::thread>> threads;
    for (int i = 0; i < datas.size(); ++i) {
      AtomicCout aout;
      aout.stream() << "Start thread " << i << std::endl;
      threads.emplace_back(std::make_unique<std::thread>(ProcessRequest, i));
    }

    // Make sure all threads have terminated.
    for (auto& thread : threads) {
      thread->join();
    }
  }
  else {
    for (size_t i = 0; i < datas.size(); ++i) {
      ProcessRequest(i);
    }
  }
  // It's safe to access global state beyond this point without locking since
  // all no more responses will be touching them.

  if (group) {
    std::cout << std::endl;
    std::cout << "[Demo] Final action for all requests is ";
    switch (global_final_action) {
    // Google Chrome fails open, so if no action is specified that is the same
    // as ALLOW.
    case ContentAnalysisAcknowledgement::ACTION_UNSPECIFIED:
    case ContentAnalysisAcknowledgement::ALLOW:
      std::cout << "allowed";
      break;
    case ContentAnalysisAcknowledgement::REPORT_ONLY:
      std::cout << "reported only";
      break;
    case ContentAnalysisAcknowledgement::WARN:
      std::cout << "warned";
      break;
    case ContentAnalysisAcknowledgement::BLOCK:
      std::cout << "blocked";
      break;
    }
    std::cout << std::endl << std::endl;

    for (auto token : request_tokens) {
      std::cout << "[Demo] Sending group Ack" << std::endl;
      int err = client->Acknowledge(
        BuildAcknowledgement(token, global_final_action));
      if (err != 0) {
        std::cout << "[Demo] Error sending ack for " << token << std::endl;
      }
    }
  }

  return 0;
};