#!/usr/bin/env bash

sc_transport_info() {
  echo "INFO: $*"
}

sc_transport_warn() {
  echo "WARN: $*" >&2
}

sc_transport_die() {
  echo "ERROR: $*" >&2
  exit 1
}

sc_transport_require_bin() {
  command -v "$1" >/dev/null 2>&1 || sc_transport_die "required command not found: $1"
}

sc_transport_init() {
  : "${SC_TRANSPORT_MODE:=ssm}"
  : "${SC_TRANSPORT_DRY_RUN:=false}"
  : "${SC_TRANSPORT_REMOTE_HOST:=}"
  : "${SC_TRANSPORT_INSTANCE_ID:=}"
  : "${SC_TRANSPORT_IDENTITY_FILE:=}"
  : "${SC_TRANSPORT_AWS_REGION:=}"
  : "${SC_TRANSPORT_AWS_PROFILE:=}"
  : "${SC_TRANSPORT_SSM_PENDING_TIMEOUT_POLLS:=60}"
  : "${SC_TRANSPORT_SSM_TOTAL_TIMEOUT_POLLS:=900}"

  case "${SC_TRANSPORT_MODE}" in
    ssm)
      sc_transport_require_bin aws
      [[ -n "${SC_TRANSPORT_INSTANCE_ID}" ]] || sc_transport_die "ssm transport requires SC_TRANSPORT_INSTANCE_ID"
      [[ -n "${SC_TRANSPORT_AWS_REGION}" ]] || sc_transport_die "ssm transport requires SC_TRANSPORT_AWS_REGION"
      export AWS_REGION="${SC_TRANSPORT_AWS_REGION}"
      export AWS_DEFAULT_REGION="${SC_TRANSPORT_AWS_REGION}"
      export AWS_SDK_LOAD_CONFIG=1
      if [[ -n "${SC_TRANSPORT_AWS_PROFILE}" ]]; then
        export AWS_PROFILE="${SC_TRANSPORT_AWS_PROFILE}"
      fi
      ;;
    ssh)
      sc_transport_require_bin ssh
      [[ -n "${SC_TRANSPORT_REMOTE_HOST}" ]] || sc_transport_die "ssh transport requires SC_TRANSPORT_REMOTE_HOST"
      ;;
    *)
      sc_transport_die "unsupported transport: ${SC_TRANSPORT_MODE}"
      ;;
  esac
}

sc_transport_target_label() {
  if [[ "${SC_TRANSPORT_MODE}" == "ssm" ]]; then
    printf '%s' "${SC_TRANSPORT_INSTANCE_ID}"
  else
    printf '%s' "${SC_TRANSPORT_REMOTE_HOST}"
  fi
}

sc_transport_ssm_get_invocation_field() {
  local command_id="$1"
  local field="$2"
  local value=""
  local attempt=""
  for attempt in 1 2 3 4 5; do
    value="$(aws ssm get-command-invocation \
      --region "${SC_TRANSPORT_AWS_REGION}" \
      --command-id "${command_id}" \
      --instance-id "${SC_TRANSPORT_INSTANCE_ID}" \
      --query "${field}" \
      --output text 2>/dev/null || true)"
    if [[ -n "${value}" && "${value}" != "None" ]]; then
      printf '%s' "${value}"
      return 0
    fi
    sleep 2
  done
  return 1
}

sc_transport_ssm_get_status() {
  local command_id="$1"
  local status=""
  status="$(sc_transport_ssm_get_invocation_field "${command_id}" 'Status' || true)"
  if [[ -n "${status}" && "${status}" != "None" ]]; then
    printf '%s' "${status}"
    return 0
  fi

  status="$(aws ssm list-command-invocations \
    --region "${SC_TRANSPORT_AWS_REGION}" \
    --command-id "${command_id}" \
    --details \
    --query 'CommandInvocations[0].Status' \
    --output text 2>/dev/null || true)"
  if [[ -n "${status}" && "${status}" != "None" ]]; then
    printf '%s' "${status}"
    return 0
  fi
  return 1
}

sc_transport_run() {
  local command="$1"

  if [[ "${SC_TRANSPORT_DRY_RUN}" == "true" ]]; then
    echo "DRY-RUN remote ($(sc_transport_target_label) via ${SC_TRANSPORT_MODE}): ${command}"
    return 0
  fi

  case "${SC_TRANSPORT_MODE}" in
    ssh)
      local -a ssh_cmd=(ssh)
      if [[ -n "${SC_TRANSPORT_IDENTITY_FILE}" ]]; then
        ssh_cmd+=(-i "${SC_TRANSPORT_IDENTITY_FILE}")
      fi
      ssh_cmd+=("${SC_TRANSPORT_REMOTE_HOST}" "${command}")
      "${ssh_cmd[@]}"
      ;;
    ssm)
      local params_file output_file stderr_file command_id status stdout_content stderr_content
      params_file="$(mktemp /tmp/security-center-ssm-params.XXXXXX)"
      output_file="$(mktemp /tmp/security-center-ssm-stdout.XXXXXX)"
      stderr_file="$(mktemp /tmp/security-center-ssm-stderr.XXXXXX)"
      python3 - "${command}" "${params_file}" <<'PY'
import json
import sys

command = sys.argv[1]
path = sys.argv[2]
payload = {"commands": [command]}
with open(path, "w", encoding="utf-8") as fh:
    json.dump(payload, fh)
PY

      command_id="$(aws ssm send-command \
        --region "${SC_TRANSPORT_AWS_REGION}" \
        --instance-ids "${SC_TRANSPORT_INSTANCE_ID}" \
        --document-name "AWS-RunShellScript" \
        --comment "security-center-enterprise-remote-command" \
        --parameters "file://${params_file}" \
        --query 'Command.CommandId' \
        --output text)"

      local poll_count=0
      while true; do
        poll_count=$((poll_count + 1))
        status="$(sc_transport_ssm_get_status "${command_id}" || true)"
        case "${status}" in
          Pending|Delayed)
            if (( poll_count == 1 || poll_count % 15 == 0 )); then
              sc_transport_info "SSM command ${command_id} status=${status}; still waiting for the instance to pick up work"
            fi
            if [[ ${poll_count} -ge ${SC_TRANSPORT_SSM_PENDING_TIMEOUT_POLLS} ]]; then
              sc_transport_warn "SSM command ${command_id} remained ${status} for too long; the managed instance may not be consuming commands"
              rm -f "${params_file}" "${output_file}" "${stderr_file}"
              return 1
            fi
            sleep 2
            ;;
          InProgress|"")
            if (( poll_count == 1 || poll_count % 30 == 0 )); then
              sc_transport_info "SSM command ${command_id} status=${status:-InProgress}; remote work is still running"
            fi
            if [[ ${poll_count} -ge ${SC_TRANSPORT_SSM_TOTAL_TIMEOUT_POLLS} ]]; then
              sc_transport_warn "timed out waiting for SSM command ${command_id}"
              rm -f "${params_file}" "${output_file}" "${stderr_file}"
              return 1
            fi
            sleep 2
            ;;
          Success)
            sc_transport_ssm_get_invocation_field "${command_id}" 'StandardOutputContent' > "${output_file}" || true
            cat "${output_file}"
            rm -f "${params_file}" "${output_file}" "${stderr_file}"
            return 0
            ;;
          Cancelled|TimedOut|Failed|Cancelling)
            sc_transport_ssm_get_invocation_field "${command_id}" 'StandardOutputContent' > "${output_file}" || true
            sc_transport_ssm_get_invocation_field "${command_id}" 'StandardErrorContent' > "${stderr_file}" || true
            stdout_content="$(cat "${output_file}" 2>/dev/null || true)"
            stderr_content="$(cat "${stderr_file}" 2>/dev/null || true)"
            [[ -n "${stdout_content}" ]] && printf '%s\n' "${stdout_content}"
            [[ -n "${stderr_content}" ]] && printf '%s\n' "${stderr_content}" >&2
            rm -f "${params_file}" "${output_file}" "${stderr_file}"
            return 1
            ;;
          *)
            sc_transport_warn "unexpected SSM command status: ${status}"
            sleep 2
            ;;
        esac
      done
      ;;
  esac
}