/**
 *
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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
 *
 *     http://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 "c2/C2Agent.h"

#include <cstdio>
#include <csignal>
#include <utility>
#include <limits>
#include <vector>
#include <map>
#include <string>
#include <memory>
#include <cinttypes>

#include "c2/ControllerSocketProtocol.h"
#include "core/ProcessContext.h"
#include "core/CoreComponentState.h"
#include "core/state/UpdateController.h"
#include "core/logging/Logger.h"
#include "core/logging/LoggerConfiguration.h"
#include "utils/file/DiffUtils.h"
#include "utils/file/FileUtils.h"
#include "utils/file/FileManager.h"
#include "utils/file/FileSystem.h"
#include "utils/GeneralUtils.h"
#include "utils/HTTPClient.h"
#include "utils/Environment.h"
#include "utils/Monitors.h"
#include "utils/OptionalUtils.h"
#include "utils/StringUtils.h"

namespace org {
namespace apache {
namespace nifi {
namespace minifi {
namespace c2 {

C2Agent::C2Agent(core::controller::ControllerServiceProvider* controller,
                 const std::shared_ptr<state::StateMonitor> &updateSink,
                 const std::shared_ptr<Configure> &configuration,
                 const std::shared_ptr<utils::file::FileSystem>& filesystem)
    : heart_beat_period_(3000),
      max_c2_responses(5),
      update_sink_(updateSink),
      update_service_(nullptr),
      controller_(controller),
      configuration_(configuration),
      filesystem_(filesystem),
      protocol_(nullptr),
      logger_(logging::LoggerFactory<C2Agent>::getLogger()),
      thread_pool_(2, false, nullptr, "C2 threadpool") {
  allow_updates_ = true;

  manifest_sent_ = false;

  running_c2_configuration = std::make_shared<Configure>();

  last_run_ = std::chrono::steady_clock::now();

  if (nullptr != controller_) {
    update_service_ = std::static_pointer_cast<controllers::UpdatePolicyControllerService>(controller_->getControllerService(UPDATE_NAME));
  }

  if (update_service_ == nullptr) {
    // create a stubbed service for updating the flow identifier
  }

  configure(configuration, false);

  functions_.emplace_back([this] {return produce();});
  functions_.emplace_back([this] {return consume();});
}

void C2Agent::start() {
  if (controller_running_) {
    return;
  }
  task_ids_.clear();
  for (const auto& function : functions_) {
    utils::Identifier uuid = utils::IdGenerator::getIdGenerator()->generate();
    task_ids_.push_back(uuid);
    auto monitor = utils::make_unique<utils::ComplexMonitor>();
    utils::Worker<utils::TaskRescheduleInfo> functor(function, uuid.to_string(), std::move(monitor));
    std::future<utils::TaskRescheduleInfo> future;
    thread_pool_.execute(std::move(functor), future);
  }
  controller_running_ = true;
  thread_pool_.start();
  logger_->log_info("C2 agent started");
}

void C2Agent::stop() {
  controller_running_ = false;
  for (const auto& id : task_ids_) {
    thread_pool_.stopTasks(id.to_string());
  }
  thread_pool_.shutdown();
  logger_->log_info("C2 agent stopped");
}

void C2Agent::checkTriggers() {
  logger_->log_debug("Checking %d triggers", triggers_.size());
  for (const auto &trigger : triggers_) {
    if (trigger->triggered()) {
      /**
       * Action was triggered, so extract it.
       */
      C2Payload &&triggerAction = trigger->getAction();
      logger_->log_trace("%s action triggered", trigger->getName());
      // handle the response the same way. This means that
      // acknowledgements will be sent to the c2 server for every trigger action.
      // this is expected
      extractPayload(std::move(triggerAction));
      // call reset if the trigger supports this activity
      trigger->reset();
    } else {
      logger_->log_trace("%s action not triggered", trigger->getName());
    }
  }
}
void C2Agent::configure(const std::shared_ptr<Configure> &configure, bool reconfigure) {
  std::string clazz, heartbeat_period, device;

  if (!reconfigure) {
    if (!configure->get("nifi.c2.agent.protocol.class", "c2.agent.protocol.class", clazz)) {
      clazz = "CoapProtocol";
    }
    logger_->log_info("Class is %s", clazz);

    auto protocol = core::ClassLoader::getDefaultClassLoader().instantiateRaw(clazz, clazz);
    if (protocol == nullptr) {
      logger_->log_warn("Class %s not found", clazz);
      protocol = core::ClassLoader::getDefaultClassLoader().instantiateRaw("CoapProtocol", "CoapProtocol");
      if (!protocol) {
        const char* errmsg = "Attempted to load CoapProtocol. To enable C2, please specify an active protocol for this agent.";
        logger_->log_error(errmsg);
        throw minifi::Exception{ minifi::GENERAL_EXCEPTION, errmsg };
      }

      logger_->log_info("Class is CoapProtocol");
    }

    // Since !reconfigure, the call comes from the ctor and protocol_ is null, therefore no delete is necessary
    protocol_.exchange(dynamic_cast<C2Protocol *>(protocol));

    protocol_.load()->initialize(controller_, configuration_);
  } else {
    protocol_.load()->update(configure);
  }

  if (configure->get("nifi.c2.agent.heartbeat.period", "c2.agent.heartbeat.period", heartbeat_period)) {
    core::TimeUnit unit;

    try {
      int64_t schedulingPeriod = 0;
      if (core::Property::StringToTime(heartbeat_period, schedulingPeriod, unit) && core::Property::ConvertTimeUnitToMS(schedulingPeriod, unit, schedulingPeriod)) {
        heart_beat_period_ = schedulingPeriod;
        logger_->log_debug("Using %u ms as the heartbeat period", heart_beat_period_);
      } else {
        heart_beat_period_ = std::stoi(heartbeat_period);
      }
    } catch (const std::invalid_argument &) {
      heart_beat_period_ = 3000;
    }
  } else {
    if (!reconfigure)
      heart_beat_period_ = 3000;
  }

  std::string update_settings;
  if (configure->get("nifi.c2.agent.update.allow", "c2.agent.update.allow", update_settings) && utils::StringUtils::StringToBool(update_settings, allow_updates_)) {
    // allow the agent to be updated. we then need to get an update command to execute after
  }

  if (allow_updates_) {
    if (!configure->get("nifi.c2.agent.update.command", "c2.agent.update.command", update_command_)) {
      std::string cwd = utils::Environment::getCurrentWorkingDirectory();
      if (cwd.empty()) {
        logger_->log_error("Could not set the update command because the working directory could not be determined");
      } else {
        update_command_ = cwd + "/minifi.sh update";
      }
    }

    if (!configure->get("nifi.c2.agent.update.temp.location", "c2.agent.update.temp.location", update_location_)) {
      std::string cwd = utils::Environment::getCurrentWorkingDirectory();
      if (cwd.empty()) {
        logger_->log_error("Could not set the update location because the working directory could not be determined");
      } else {
        update_location_ = cwd + "/minifi.update";
      }
    }

    // if not defined we won't be able to update
    configure->get("nifi.c2.agent.bin.location", "c2.agent.bin.location", bin_location_);
  }
  std::string heartbeat_reporters;
  if (configure->get("nifi.c2.agent.heartbeat.reporter.classes", "c2.agent.heartbeat.reporter.classes", heartbeat_reporters)) {
    std::vector<std::string> reporters = utils::StringUtils::split(heartbeat_reporters, ",");
    std::lock_guard<std::mutex> lock(heartbeat_mutex);
    for (const auto& reporter : reporters) {
      auto heartbeat_reporter_obj = core::ClassLoader::getDefaultClassLoader().instantiate(reporter, reporter);
      if (heartbeat_reporter_obj == nullptr) {
        logger_->log_debug("Could not instantiate %s", reporter);
      } else {
        std::shared_ptr<HeartBeatReporter> shp_reporter = std::static_pointer_cast<HeartBeatReporter>(heartbeat_reporter_obj);
        shp_reporter->initialize(controller_, update_sink_, configuration_);
        heartbeat_protocols_.push_back(shp_reporter);
      }
    }
  }

  std::string trigger_classes;
  if (configure->get("nifi.c2.agent.trigger.classes", "c2.agent.trigger.classes", trigger_classes)) {
    std::vector<std::string> triggers = utils::StringUtils::split(trigger_classes, ",");
    std::lock_guard<std::mutex> lock(heartbeat_mutex);
    for (const auto& trigger : triggers) {
      auto trigger_obj = core::ClassLoader::getDefaultClassLoader().instantiate(trigger, trigger);
      if (trigger_obj == nullptr) {
        logger_->log_debug("Could not instantiate %s", trigger);
      } else {
        std::shared_ptr<C2Trigger> trg_impl = std::static_pointer_cast<C2Trigger>(trigger_obj);
        trg_impl->initialize(configuration_);
        triggers_.push_back(trg_impl);
      }
    }
  }

  auto base_reporter = "ControllerSocketProtocol";
  auto heartbeat_reporter_obj = core::ClassLoader::getDefaultClassLoader().instantiate(base_reporter, base_reporter);
  if (heartbeat_reporter_obj == nullptr) {
    logger_->log_debug("Could not instantiate %s", base_reporter);
  } else {
    std::shared_ptr<HeartBeatReporter> shp_reporter = std::static_pointer_cast<HeartBeatReporter>(heartbeat_reporter_obj);
    shp_reporter->initialize(controller_, update_sink_, configuration_);
    heartbeat_protocols_.push_back(shp_reporter);
  }
}

void C2Agent::performHeartBeat() {
  C2Payload payload(Operation::HEARTBEAT);
  logger_->log_trace("Performing heartbeat");
  std::shared_ptr<state::response::NodeReporter> reporter = std::dynamic_pointer_cast<state::response::NodeReporter>(update_sink_);
  std::vector<std::shared_ptr<state::response::ResponseNode>> metrics;
  if (reporter) {
    if (!manifest_sent_) {
      // include agent manifest for the first heartbeat
      metrics = reporter->getHeartbeatNodes(true);
      manifest_sent_ = true;
    } else {
      metrics = reporter->getHeartbeatNodes(false);
    }

    payload.reservePayloads(metrics.size());
    for (const auto& metric : metrics) {
      C2Payload child_metric_payload(Operation::HEARTBEAT);
      child_metric_payload.setLabel(metric->getName());
      child_metric_payload.setContainer(metric->isArray());
      serializeMetrics(child_metric_payload, metric->getName(), metric->serialize(), metric->isArray());
      payload.addPayload(std::move(child_metric_payload));
    }
  }
  C2Payload && response = protocol_.load()->consumePayload(payload);

  enqueue_c2_server_response(std::move(response));

  std::lock_guard<std::mutex> lock(heartbeat_mutex);

  for (auto reporter : heartbeat_protocols_) {
    reporter->heartbeat(payload);
  }
}

void C2Agent::serializeMetrics(C2Payload &metric_payload, const std::string &name, const std::vector<state::response::SerializedResponseNode> &metrics, bool is_container, bool is_collapsible) {
  const auto payloads = std::count_if(begin(metrics), end(metrics), [](const state::response::SerializedResponseNode& metric) { return !metric.children.empty(); });
  metric_payload.reservePayloads(metric_payload.getNestedPayloads().size() + payloads);
  for (const auto &metric : metrics) {
    if (metric.children.size() > 0) {
      C2Payload child_metric_payload(metric_payload.getOperation());
      if (metric.array) {
        child_metric_payload.setContainer(true);
      }
      auto collapsible = !metric.collapsible ? metric.collapsible : is_collapsible;
      child_metric_payload.setCollapsible(collapsible);
      child_metric_payload.setLabel(metric.name);
      serializeMetrics(child_metric_payload, metric.name, metric.children, is_container, collapsible);
      metric_payload.addPayload(std::move(child_metric_payload));
    } else {
      C2ContentResponse response(metric_payload.getOperation());
      response.name = name;
      response.operation_arguments[metric.name] = metric.value;
      metric_payload.addContent(std::move(response), is_collapsible);
    }
  }
}

void C2Agent::extractPayload(const C2Payload &resp) {
  if (resp.getStatus().getState() == state::UpdateState::NESTED) {
    for (const C2Payload& payload : resp.getNestedPayloads()) {
      extractPayload(payload);
    }
    return;
  }
  switch (resp.getStatus().getState()) {
    case state::UpdateState::INITIATE:
      logger_->log_debug("Received initiation event from protocol");
      break;
    case state::UpdateState::READ_COMPLETE:
      logger_->log_trace("Received Ack from Server");
      // we have a heartbeat response.
      for (const C2ContentResponse &server_response : resp.getContent()) {
        handle_c2_server_response(server_response);
      }
      break;
    case state::UpdateState::FULLY_APPLIED:
      logger_->log_debug("Received fully applied event from protocol");
      break;
    case state::UpdateState::PARTIALLY_APPLIED:
      logger_->log_debug("Received partially applied event from protocol");
      break;
    case state::UpdateState::NOT_APPLIED:
      logger_->log_debug("Received not applied event from protocol");
      break;
    case state::UpdateState::SET_ERROR:
      logger_->log_debug("Received error event from protocol");
      break;
    case state::UpdateState::READ_ERROR:
      logger_->log_debug("Received error event from protocol");
      break;
    case state::UpdateState::NESTED:  // multiple updates embedded into one

    default:
      logger_->log_debug("Received nested event from protocol");
      break;
  }
}

void C2Agent::handle_c2_server_response(const C2ContentResponse &resp) {
  switch (resp.op) {
    case Operation::CLEAR:
      // we've been told to clear something
      if (resp.name == "connection") {
        for (const auto& connection : resp.operation_arguments) {
          logger_->log_debug("Clearing connection %s", connection.second.to_string());
          update_sink_->clearConnection(connection.second.to_string());
        }
        C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
        enqueue_c2_response(std::move(response));
      } else if (resp.name == "repositories") {
        update_sink_->drainRepositories();
        C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
        enqueue_c2_response(std::move(response));
      } else if (resp.name == "corecomponentstate") {
        // TODO(bakaid): untested
        std::vector<std::shared_ptr<state::StateController>> components = update_sink_->getComponents(resp.name);
        auto state_manager_provider = core::ProcessContext::getStateManagerProvider(logger_, controller_, configuration_);
        if (state_manager_provider != nullptr) {
          for (auto &component : components) {
            logger_->log_debug("Clearing state for component %s", component->getComponentName());
            auto state_manager = state_manager_provider->getCoreComponentStateManager(component->getComponentUUID());
            if (state_manager != nullptr) {
              component->stop();
              state_manager->clear();
              state_manager->persist();
              component->start();
            } else {
              logger_->log_warn("Failed to get StateManager for component %s", component->getComponentUUID().to_string());
            }
          }
        } else {
          logger_->log_error("Failed to get StateManagerProvider");
        }
        C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
        enqueue_c2_response(std::move(response));
      } else {
        logger_->log_debug("Clearing unknown %s", resp.name);
      }

      break;
    case Operation::UPDATE:
      handle_update(resp);
      break;
    case Operation::DESCRIBE:
      handle_describe(resp);
      break;
    case Operation::RESTART: {
      update_sink_->stop();
      C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
      protocol_.load()->consumePayload(std::move(response));
      restart_agent();
    }
      break;
    case Operation::START:
    case Operation::STOP: {
      if (resp.name == "C2" || resp.name == "c2") {
        raise(SIGTERM);
      }

      std::vector<std::shared_ptr<state::StateController>> components = update_sink_->getComponents(resp.name);
      // stop all referenced components.
      for (auto &component : components) {
        logger_->log_debug("Stopping component %s", component->getComponentName());
        if (resp.op == Operation::STOP)
          component->stop();
        else
          component->start();
      }

      if (resp.ident.length() > 0) {
        C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
        enqueue_c2_response(std::move(response));
      }
    }
      //
      break;
    default:
      break;
      // do nothing
  }
}

C2Payload C2Agent::prepareConfigurationOptions(const C2ContentResponse &resp) const {
    auto unsanitized_keys = configuration_->getConfiguredKeys();
    std::vector<std::string> keys;
    std::copy_if(unsanitized_keys.begin(), unsanitized_keys.end(), std::back_inserter(keys),
            [](std::string key) {return key.find("pass") == std::string::npos;});

    C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
    C2Payload options(Operation::ACKNOWLEDGE);
    options.setLabel("configuration_options");
    std::string value;
    for (auto key : keys) {
      C2ContentResponse option(Operation::ACKNOWLEDGE);
      option.name = key;
      if (configuration_->get(key, value)) {
        option.operation_arguments[key] = value;
        options.addContent(std::move(option));
      }
    }
    response.addPayload(std::move(options));
    return response;
}

/**
 * Descriptions are special types of requests that require information
 * to be put into the acknowledgement
 */
void C2Agent::handle_describe(const C2ContentResponse &resp) {
  auto reporter = std::dynamic_pointer_cast<state::response::NodeReporter>(update_sink_);
  if (resp.name == "metrics") {
    C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
    if (reporter != nullptr) {
      auto iter = resp.operation_arguments.find("metricsClass");
      std::string metricsClass;
      if (iter != resp.operation_arguments.end()) {
        metricsClass = iter->second.to_string();
      }
      auto metricsNode = reporter->getMetricsNode(metricsClass);
      C2Payload metrics(Operation::ACKNOWLEDGE);
      metricsClass.empty() ? metrics.setLabel("metrics") : metrics.setLabel(metricsClass);
      if (metricsNode) {
        serializeMetrics(metrics, metricsNode->getName(), metricsNode->serialize(), metricsNode->isArray());
      }
      response.addPayload(std::move(metrics));
    }
    enqueue_c2_response(std::move(response));
    return;
  } else if (resp.name == "configuration") {
    auto configOptions = prepareConfigurationOptions(resp);
    enqueue_c2_response(std::move(configOptions));
    return;
  } else if (resp.name == "manifest") {
    C2Payload response(prepareConfigurationOptions(resp));
    if (reporter != nullptr) {
      C2Payload agentInfo(Operation::ACKNOWLEDGE, resp.ident, true);
      agentInfo.setLabel("agentInfo");

      const auto manifest = reporter->getAgentManifest();
      serializeMetrics(agentInfo, manifest->getName(), manifest->serialize());
      response.addPayload(std::move(agentInfo));
    }
    enqueue_c2_response(std::move(response));
    return;
  } else if (resp.name == "jstack") {
    if (update_sink_->isRunning()) {
      const std::vector<BackTrace> traces = update_sink_->getTraces();
      for (const auto &trace : traces) {
        for (const auto & line : trace.getTraces()) {
          logger_->log_trace("%s -- %s", trace.getName(), line);
        }
      }
      auto keys = configuration_->getConfiguredKeys();
      C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
      for (const auto &trace : traces) {
        C2Payload options(Operation::ACKNOWLEDGE, resp.ident, true);
        options.setLabel(trace.getName());
        std::string value;
        for (const auto &line : trace.getTraces()) {
          C2ContentResponse option(Operation::ACKNOWLEDGE);
          option.name = line;
          option.operation_arguments[line] = line;
          options.addContent(std::move(option));
        }
        response.addPayload(std::move(options));
      }
      enqueue_c2_response(std::move(response));
      return;
    }
  } else if (resp.name == "corecomponentstate") {
    C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
    response.setLabel("corecomponentstate");
    C2Payload states(Operation::ACKNOWLEDGE, resp.ident, true);
    states.setLabel("corecomponentstate");
    auto state_manager_provider = core::ProcessContext::getStateManagerProvider(logger_, controller_, configuration_);
    if (state_manager_provider != nullptr) {
      auto core_component_states = state_manager_provider->getAllCoreComponentStates();
      for (const auto& core_component_state : core_component_states) {
        C2Payload state(Operation::ACKNOWLEDGE, resp.ident, true);
        state.setLabel(core_component_state.first.to_string());
        for (const auto& kv : core_component_state.second) {
          C2ContentResponse entry(Operation::ACKNOWLEDGE);
          entry.name = kv.first;
          entry.operation_arguments[kv.first] = kv.second;
          state.addContent(std::move(entry));
        }
        states.addPayload(std::move(state));
      }
    }
    response.addPayload(std::move(states));
    enqueue_c2_response(std::move(response));
    return;
  }
  C2Payload response(Operation::ACKNOWLEDGE, resp.ident, true);
  enqueue_c2_response(std::move(response));
}

void C2Agent::handle_update(const C2ContentResponse &resp) {
  // we've been told to update something
  if (resp.name == "configuration") {
    handleConfigurationUpdate(resp);
  } else if (resp.name == "properties") {
    state::UpdateState result = state::UpdateState::FULLY_APPLIED;
    for (auto entry : resp.operation_arguments) {
      bool persist = (
          entry.second.getAnnotation("persist")
          | utils::map(&AnnotatedValue::to_string)
          | utils::flatMap(utils::StringUtils::toBool)).value_or(false);
      if (!update_property(entry.first, entry.second.to_string(), persist)) {
        result = state::UpdateState::PARTIALLY_APPLIED;
      }
    }
    C2Payload response(Operation::ACKNOWLEDGE, result, resp.ident, true);
    enqueue_c2_response(std::move(response));
  } else if (resp.name == "c2") {
    // prior configuration options were already in place. thus
    // we clear the map so that we don't go through replacing
    // unnecessary objects.
    running_c2_configuration->clear();

    for (auto entry : resp.operation_arguments) {
      bool can_update = true;
      if (nullptr != update_service_) {
        can_update = update_service_->canUpdate(entry.first);
      }
      if (can_update)
        running_c2_configuration->set(entry.first, entry.second.to_string());
    }

    if (resp.operation_arguments.size() > 0)
      configure(running_c2_configuration);
    C2Payload response(Operation::ACKNOWLEDGE, state::UpdateState::FULLY_APPLIED, resp.ident, true);
    enqueue_c2_response(std::move(response));
  } else if (resp.name == "agent") {
    // we are upgrading the agent. therefore we must be given a location
    auto location = resp.operation_arguments.find("location");
    auto isPartialStr = resp.operation_arguments.find("partial");

    bool partial_update = false;
    if (isPartialStr != std::end(resp.operation_arguments)) {
      partial_update = utils::StringUtils::equalsIgnoreCase(isPartialStr->second.to_string(), "true");
    }
    if (location != resp.operation_arguments.end()) {
      logger_->log_trace("Update agent with location %s", location->second.to_string());
      // we will not have a raw payload
      C2Payload payload(Operation::TRANSFER, true);

      C2Payload &&response = protocol_.load()->consumePayload(location->second.to_string(), payload, RECEIVE, false);

      auto raw_data = response.getRawData();

      std::string file_path = std::string(raw_data.data(), raw_data.size());

      logger_->log_trace("Update requested with file %s", file_path);

      // acknowledge the transfer. For a transfer, the response identifier should be the checksum of the
      // file transferred.
      C2Payload transfer_response(Operation::ACKNOWLEDGE, state::UpdateState::FULLY_APPLIED, response.getIdentifier(), true);

      protocol_.load()->consumePayload(std::move(transfer_response));

      if (allow_updates_) {
        logger_->log_trace("Update allowed from file %s", file_path);
        if (partial_update && !bin_location_.empty()) {
          utils::file::DiffUtils::apply_binary_diff(bin_location_.c_str(), file_path.c_str(), update_location_.c_str());
        } else {
          utils::file::FileUtils::copy_file(file_path, update_location_);
        }
        // remove the downloaded file.
        logger_->log_trace("removing file %s", file_path);
        std::remove(file_path.c_str());
        update_agent();
      } else {
        logger_->log_trace("Update disallowed from file %s", file_path);
      }

    } else {
      logger_->log_trace("No location present");
    }
  } else {
    C2Payload response(Operation::ACKNOWLEDGE, state::UpdateState::NOT_APPLIED, resp.ident, true);
    enqueue_c2_response(std::move(response));
  }
}

/**
 * Updates a property
 */
bool C2Agent::update_property(const std::string &property_name, const std::string &property_value, bool persist) {
  if (update_service_ && !update_service_->canUpdate(property_name)) {
    return false;
  }
  configuration_->set(property_name, property_value);
  if (!persist) {
    return true;
  }
  return configuration_->persistProperties();
}

void C2Agent::restart_agent() {
  std::string cwd = utils::Environment::getCurrentWorkingDirectory();
  if (cwd.empty()) {
    logger_->log_error("Could not restart the agent because the working directory could not be determined");
    return;
  }

  std::string command = cwd + "/bin/minifi.sh restart";
  if (system(command.c_str()) != 0) {
    logger_->log_warn("System command '%s' failed", command);
  }
}

void C2Agent::update_agent() {
  if (system(update_command_.c_str()) == 0) {
    logger_->log_info("Executed update command '%s'", update_command_);
  } else {
    logger_->log_warn("Update command %s failed; may not have command processor", update_command_);
  }
}

utils::TaskRescheduleInfo C2Agent::produce() {
  // place priority on messages to send to the c2 server
  if (protocol_.load() != nullptr) {
    std::vector<C2Payload> payload_batch;
    payload_batch.reserve(max_c2_responses);
    auto getRequestPayload = [&payload_batch] (C2Payload&& payload) { payload_batch.emplace_back(std::move(payload)); };
    for (std::size_t attempt_num = 0; attempt_num < max_c2_responses; ++attempt_num) {
      if (!requests.consume(getRequestPayload)) {
        break;
      }
    }
    std::for_each(
        std::make_move_iterator(payload_batch.begin()),
        std::make_move_iterator(payload_batch.end()),
        [&] (C2Payload&& payload) {
          try {
            C2Payload && response = protocol_.load()->consumePayload(std::move(payload));
            enqueue_c2_server_response(std::move(response));
          }
          catch(const std::exception &e) {
            logger_->log_error("Exception occurred while consuming payload. error: %s", e.what());
          }
          catch(...) {
            logger_->log_error("Unknown exception occurred while consuming payload.");
          }
        });

    try {
      performHeartBeat();
    }
    catch (const std::exception &e) {
      logger_->log_error("Exception occurred while performing heartbeat. error: %s", e.what());
    }
    catch (...) {
      logger_->log_error("Unknonwn exception occurred while performing heartbeat.");
    }
  }

  checkTriggers();

  return utils::TaskRescheduleInfo::RetryIn(std::chrono::milliseconds(heart_beat_period_));
}

utils::TaskRescheduleInfo C2Agent::consume() {
  const auto consume_success = responses.consume([this] (C2Payload&& payload) {
    extractPayload(std::move(payload));
  });
  if (!consume_success) {
    extractPayload(C2Payload{ Operation::HEARTBEAT });
  }
  return utils::TaskRescheduleInfo::RetryIn(std::chrono::milliseconds(C2RESPONSE_POLL_MS));
}

utils::optional<std::string> C2Agent::fetchFlow(const std::string& uri) const {
  if (!utils::StringUtils::startsWith(uri, "http") || protocol_.load() == nullptr) {
    // try to open the file
    utils::optional<std::string> content = filesystem_->read(uri);
    if (content) {
      return content;
    }
  }
  // couldn't open as file and we have no protocol to request the file from
  if (protocol_.load() == nullptr) {
    return {};
  }

  std::string resolved_url = uri;
  if (!utils::StringUtils::startsWith(uri, "http")) {
    std::stringstream adjusted_url;
    std::string base;
    if (configuration_->get(minifi::Configure::nifi_c2_flow_base_url, base)) {
      adjusted_url << base;
      if (!utils::StringUtils::endsWith(base, "/")) {
        adjusted_url << "/";
      }
      adjusted_url << uri;
      resolved_url = adjusted_url.str();
    } else if (configuration_->get("c2.rest.url", base)) {
      std::string host, protocol;
      int port = -1;
      utils::parse_url(&base, &host, &port, &protocol);
      adjusted_url << protocol << host;
      if (port > 0) {
        adjusted_url << ":" << port;
      }
      adjusted_url << "/c2/api/" << uri;
      resolved_url = adjusted_url.str();
    }
  }

  C2Payload payload(Operation::TRANSFER, true);
  C2Payload &&response = protocol_.load()->consumePayload(resolved_url, payload, RECEIVE, false);

  auto raw_data = response.getRawData();
  return std::string(raw_data.data(), raw_data.size());
}

bool C2Agent::handleConfigurationUpdate(const C2ContentResponse &resp) {
  auto url = resp.operation_arguments.find("location");

  std::string file_uri;
  std::string configuration_str;

  if (url != resp.operation_arguments.end()) {
    file_uri = url->second.to_string();
    utils::optional<std::string> optional_configuration_str = fetchFlow(file_uri);
    if (!optional_configuration_str) {
      logger_->log_debug("Couldn't load new flow configuration from: \"%s\"", file_uri);
      C2Payload response(Operation::ACKNOWLEDGE, state::UpdateState::SET_ERROR, resp.ident, true);
      response.setRawData("Error while applying flow. Couldn't load flow configuration.");
      enqueue_c2_response(std::move(response));
      return false;
    }
    configuration_str = optional_configuration_str.value();
  } else {
    logger_->log_debug("Did not have location within %s", resp.ident);
    auto update_text = resp.operation_arguments.find("configuration_data");
    if (update_text == resp.operation_arguments.end()) {
      logger_->log_debug("Neither the config file location nor the data is provided");
      C2Payload response(Operation::ACKNOWLEDGE, state::UpdateState::SET_ERROR, resp.ident, true);
      response.setRawData("Error while applying flow. Neither the config file location nor the data is provided.");
      enqueue_c2_response(std::move(response));
      return false;
    }
    configuration_str = update_text->second.to_string();
  }

  bool should_persist = [&] {
    auto persist = resp.operation_arguments.find("persist");
    if (persist == resp.operation_arguments.end()) {
      return false;
    }
    return utils::StringUtils::equalsIgnoreCase(persist->second.to_string(), "true");
  }();

  int16_t err = {update_sink_->applyUpdate(file_uri, configuration_str, should_persist)};
  if (err != 0) {
    logger_->log_debug("Flow configuration update failed with error code %" PRIi16, err);
    C2Payload response(Operation::ACKNOWLEDGE, state::UpdateState::SET_ERROR, resp.ident, true);
    response.setRawData("Error while applying flow. Likely missing processors");
    enqueue_c2_response(std::move(response));
    return false;
  }

  C2Payload response(Operation::ACKNOWLEDGE, state::UpdateState::FULLY_APPLIED, resp.ident, true);
  enqueue_c2_response(std::move(response));

  if (should_persist) {
    // update the flow id
    configuration_->persistProperties();
  }

  return true;
}

}  // namespace c2
}  // namespace minifi
}  // namespace nifi
}  // namespace apache
}  // namespace org