raven_integration_mongo_db

defmodule Integrations.MongoDb do
  @moduledoc """
  MongoDB server monitor.

  Collection only — see `Integrations.MongoDb.Display` (same package) for
  the dashboard panel. `Display.BundledDefault` auto-hooks it whenever this
  monitor starts, same as a single-module package would; a release without
  `raven_web` simply never compiles the display half and runs this monitor
  headless.

  Implements the MongoDB wire protocol (OP_MSG, opcode 2013) and BSON encoding
  over raw TCP — no external library required. Authenticates with SCRAM-SHA-256
  using only OTP's `:crypto` module (requires OTP 24+).

  Collects connection counts, operation throughput, memory usage, and replica
  set role from `hello` and `serverStatus`. Optionally collects collection and
  size stats from `dbStats` when a database name is provided.

  ## Params

    * `:host`                — MongoDB hostname or IP. Required.
    * `:port`                — Port number. Defaults to `27017`.
    * `:user`                — Username. Leave blank to skip authentication.
    * `:password`            — Password.
    * `:database`            — Database for `dbStats` (optional).
    * `:auth_db`             — Authentication database. Defaults to `admin`.
    * `:timeout_ms`          — TCP connect and command timeout. Defaults to `5000`.
    * `:latency_degraded_ms` — `hello` round-trip threshold for `:degraded`. Defaults to `500`.

  ## Health signal

    * `:up`       — Connected, authenticated, latency within threshold.
    * `:degraded` — Latency above threshold.
    * `:down`     — Connection refused, timeout, or authentication failure.

  ## Metrics

    * `latency_ms`       — `hello` round-trip time
    * `connections`      — Current connection count
    * `ops_rate`         — Total operations per second (rolling delta)
    * `query_rate`       — Queries per second
    * `insert_rate`      — Inserts per second
    * `update_rate`      — Updates per second
    * `delete_rate`      — Deletes per second
    * `mem_resident_mb`  — Resident memory in MB
  """

  use CodeNameRaven.Monitor

  @default_port            27017
  @default_timeout_ms      5_000
  @default_lat_degraded_ms 500

  @impl true
  def params_template do
    %{
      host:     "mongodb.example.com",
      port:     "27017",
      user:     "root",
      password: "",
      database: "",
      auth_db:  "admin"
    }
  end

  @impl true
  def params_schema do
    [
      host:                [type: :string,          required: true,  doc: "MongoDB hostname or IP"],
      port:                [type: :non_neg_integer, default: 27017,  doc: "Port number"],
      user:                [type: :string,          required: false, doc: "Username (blank = no auth)"],
      password:            [type: :string,          required: false, doc: "Password"],
      database:            [type: :string,          required: false, doc: "Database for dbStats (optional)"],
      auth_db:             [type: :string,          default: "admin", doc: "Authentication database"],
      timeout_ms:          [type: :non_neg_integer, default: 5_000,  doc: "TCP connect and command timeout in milliseconds"],
      latency_degraded_ms: [type: :non_neg_integer, default: 500,    doc: "hello round-trip threshold for degraded"]
    ]
  end

  @impl true
  def target_uri(params) do
    host = get_param(params, :host)
    port = parse_int(get_param(params, :port), @default_port)
    user = get_param(params, :user)
    db   = get_param(params, :database) || get_param(params, :auth_db) || "admin"
    cond do
      is_nil(host) -> :none
      user -> {:ok, "mongodb://#{user}@#{host}:#{port}/#{db}"}
      true -> {:ok, "mongodb://#{host}:#{port}"}
    end
  end

  @impl true
  def identity_params(params) do
    %{
      host:     get_param(params, :host),
      port:     parse_int(get_param(params, :port), @default_port),
      user:     get_param(params, :user),
      database: get_param(params, :database)
    }
  end

  # ---------------------------------------------------------------------------
  # Collect
  # ---------------------------------------------------------------------------

  @impl true
  def collect(params, state) do
    host     = get_param(params, :host)
    port     = parse_int(get_param(params, :port), @default_port)
    user     = get_param(params, :user)
    password = get_param(params, :password) || ""
    database = get_param(params, :database)
    auth_db  = get_param(params, :auth_db) || "admin"
    timeout  = parse_int(get_param(params, :timeout_ms), @default_timeout_ms)

    if is_nil(host) do
      {:error, "missing required param :host", state}
    else
      tcp_opts = [:binary, active: false, packet: :raw, send_timeout: timeout]

      case :gen_tcp.connect(to_charlist(host), port, tcp_opts, timeout) do
        {:ok, socket} ->
          result = run_checks(socket, user, password, auth_db, database, timeout, state)
          :gen_tcp.close(socket)
          case result do
            {:ok, data, new_state} -> {:ok, data, new_state}
            {:error, reason}       -> {:error, reason, state}
          end

        {:error, :econnrefused} -> {:error, "connection refused on port #{port}", state}
        {:error, :timeout}      -> {:error, "connection timed out after #{timeout}ms", state}
        {:error, :nxdomain}     -> {:error, "hostname not found: #{host}", state}
        {:error, reason}        -> {:error, inspect(reason), state}
      end
    end
  end

  # ---------------------------------------------------------------------------
  # Healthy?
  # ---------------------------------------------------------------------------

  @impl true
  def healthy?(%{assertions: [_ | _], assertion_status: status}), do: status

  def healthy?(%{latency_ms: lat})
      when is_integer(lat) and lat >= @default_lat_degraded_ms,
      do: :degraded

  def healthy?(_), do: :up

  # ---------------------------------------------------------------------------
  # Metrics
  # ---------------------------------------------------------------------------

  @impl true
  def metrics(result) do
    ops = result[:ops] || %{}

    %{latency_ms: result.latency_ms}
    |> maybe_put(:connections,     get_in(result, [:connections, :current]))
    |> maybe_put(:ops_rate,        ops[:total_rate])
    |> maybe_put(:query_rate,      ops[:query_rate])
    |> maybe_put(:insert_rate,     ops[:insert_rate])
    |> maybe_put(:update_rate,     ops[:update_rate])
    |> maybe_put(:delete_rate,     ops[:delete_rate])
    |> maybe_put(:mem_resident_mb, get_in(result, [:memory, :resident_mb]))
  end

  # ---------------------------------------------------------------------------
  # Check orchestration
  # ---------------------------------------------------------------------------

  defp run_checks(socket, user, password, auth_db, database, timeout, state) do
    t0 = System.monotonic_time(:millisecond)

    with {:ok, hello} <- command(socket, 1, "admin", [{"hello", 1}], timeout),
         :ok          <- check_ok(hello) do

      latency_ms = System.monotonic_time(:millisecond) - t0

      with :ok          <- maybe_auth(socket, user, password, auth_db, timeout),
           {:ok, status} <- command(socket, 10, "admin", [
             {"serverStatus",              1},
             {"wiredTiger",                0},
             {"tcmalloc",                  0},
             {"opLatencies",               0},
             {"logicalSessionRecordCache", 0}
           ], timeout),
           :ok <- check_ok(status) do

        db_stats = collect_db_stats(socket, database, timeout)

        now      = DateTime.utc_now()
        elapsed  = elapsed_sec(state[:prev_collected_at], now)
        prev_ops = state[:prev_opcounters] || %{}

        opcounters = status["opcounters"] || %{}
        conns      = status["connections"] || %{}
        mem        = status["mem"] || %{}

        is_primary   = hello["isWritablePrimary"] || hello["ismaster"] || false
        is_secondary = hello["secondary"] || false
        replica_set  = hello["setName"]
        version      = status["version"]

        result = %{
          latency_ms:   latency_ms,
          version:      version,
          is_primary:   is_primary,
          is_secondary: is_secondary,
          replica_set:  replica_set,
          connections: %{
            current:   int_val(conns["current"]),
            available: int_val(conns["available"])
          },
          ops: %{
            total_rate:  total_ops_rate(opcounters, prev_ops, elapsed),
            insert_rate: op_rate(opcounters, prev_ops, "insert", elapsed),
            query_rate:  op_rate(opcounters, prev_ops, "query",  elapsed),
            update_rate: op_rate(opcounters, prev_ops, "update", elapsed),
            delete_rate: op_rate(opcounters, prev_ops, "delete", elapsed)
          },
          memory: %{
            resident_mb: int_val(mem["resident"]),
            virtual_mb:  int_val(mem["virtual"])
          },
          db_stats: db_stats
        }

        new_state = %{prev_opcounters: opcounters, prev_collected_at: now}
        {:ok, result, new_state}
      end
    end
  end

  defp collect_db_stats(_socket, nil, _timeout), do: nil
  defp collect_db_stats(socket, database, timeout) do
    case command(socket, 20, database, [{"dbStats", 1}], timeout) do
      {:ok, s} when is_map(s) ->
        %{
          collections:        int_val(s["collections"]),
          objects:            int_val(s["objects"]),
          data_size_bytes:    num_val(s["dataSize"]),
          storage_size_bytes: num_val(s["storageSize"]),
          index_size_bytes:   num_val(s["indexSize"])
        }
      _ -> nil
    end
  end

  # ---------------------------------------------------------------------------
  # SCRAM-SHA-256 authentication
  # ---------------------------------------------------------------------------

  defp maybe_auth(_socket, nil, _pw, _db, _timeout), do: :ok
  defp maybe_auth(_socket, "",  _pw, _db, _timeout), do: :ok
  defp maybe_auth(socket, user, password, auth_db, timeout) do
    cnonce  = :base64.encode(:crypto.strong_rand_bytes(24))
    cfm_bare = "n=#{sasl_escape(user)},r=#{cnonce}"
    cfm      = "n,," <> cfm_bare

    with {:ok, r1} <- command(socket, 2, auth_db, [
           {"saslStart", 1},
           {"mechanism", "SCRAM-SHA-256"},
           {"payload",   {:bin, cfm}},
           {"options",   %{"skipEmptyExchange" => true}}
         ], timeout),
         :ok <- check_ok(r1),
         {:ok, snonce, salt, iters} <- parse_server_first(bin_payload(r1["payload"])),
         true <- String.starts_with?(snonce, cnonce) || {:error, "SCRAM nonce mismatch"} do

      conv_id     = r1["conversationId"]
      salted_pw   = :crypto.pbkdf2_hmac(:sha256, password, salt, iters, 32)
      client_key  = :crypto.mac(:hmac, :sha256, salted_pw, "Client Key")
      stored_key  = :crypto.hash(:sha256, client_key)
      cfmwp       = "c=biws,r=#{snonce}"
      auth_msg    = cfm_bare <> "," <> bin_payload(r1["payload"]) <> "," <> cfmwp
      client_sig  = :crypto.mac(:hmac, :sha256, stored_key, auth_msg)
      client_proof = :crypto.exor(client_key, client_sig)
      cfinal      = cfmwp <> ",p=" <> :base64.encode(client_proof)

      with {:ok, r2} <- command(socket, 3, auth_db, [
             {"saslContinue",   1},
             {"conversationId", conv_id},
             {"payload",        {:bin, cfinal}}
           ], timeout),
           :ok <- check_ok(r2) do
        :ok
      end
    else
      false        -> {:error, "SCRAM nonce mismatch"}
      {:error, r}  -> {:error, r}
    end
  end

  defp parse_server_first(sfm) when is_binary(sfm) do
    parts =
      sfm
      |> String.split(",")
      |> Map.new(fn part ->
        case String.split(part, "=", parts: 2) do
          [k, v] -> {k, v}
          _      -> {"", ""}
        end
      end)

    with r when is_binary(r) <- parts["r"],
         s when is_binary(s) <- parts["s"],
         i when is_binary(i) <- parts["i"],
         {iters, _}          <- Integer.parse(i) do
      {:ok, r, :base64.decode(s), iters}
    else
      _ -> {:error, "SCRAM: could not parse server-first-message: #{sfm}"}
    end
  end
  defp parse_server_first(_), do: {:error, "SCRAM: missing or invalid server payload"}

  defp bin_payload({:binary, _subtype, data}), do: data
  defp bin_payload(other) when is_binary(other), do: other
  defp bin_payload(_), do: ""

  defp sasl_escape(user) do
    user
    |> String.replace("=", "=3D")
    |> String.replace(",", "=2C")
  end

  # ---------------------------------------------------------------------------
  # MongoDB OP_MSG wire protocol
  # ---------------------------------------------------------------------------

  # `doc` must be an ordered list of {key, value} pairs, command name first —
  # see encode_doc/1's moduledoc note. $db is appended last, matching the
  # driver convention (and irrelevant to command dispatch either way).
  defp command(socket, req_id, db, doc, timeout) do
    bson     = encode_doc(doc ++ [{"$db", db}])
    body_len = 4 + 1 + byte_size(bson)
    total    = 16 + body_len

    msg = <<total::little-32,
             req_id::little-32,
             0::little-32,
             2013::little-32,
             0::little-32,
             0::8,
             bson::binary>>

    with :ok <- :gen_tcp.send(socket, msg) do
      recv_msg(socket, timeout)
    end
  end

  defp recv_msg(socket, timeout) do
    case :gen_tcp.recv(socket, 16, timeout) do
      {:ok, <<total::little-32, _req::little-32, _resp::little-32, _op::little-32>>} ->
        body_len = total - 16
        case :gen_tcp.recv(socket, body_len, timeout) do
          {:ok, <<_flags::little-32, 0, bson::binary>>} ->
            {:ok, decode_doc(bson)}
          {:ok, _other} ->
            {:error, "unexpected OP_MSG section kind"}
          {:error, r} ->
            {:error, "recv body failed: #{inspect(r)}"}
        end
      {:error, r} ->
        {:error, "recv header failed: #{inspect(r)}"}
    end
  end

  defp check_ok(%{"ok" => v}) when v in [1, 1.0], do: :ok
  defp check_ok(%{"errmsg" => msg}), do: {:error, "MongoDB error: #{msg}"}
  defp check_ok(_), do: {:error, "MongoDB command failed (ok != 1)"}

  # ---------------------------------------------------------------------------
  # BSON encoder
  # ---------------------------------------------------------------------------

  # A command document's first field is how the MongoDB server decides which
  # command is being invoked — field order matters here, and a Map's
  # iteration order is not guaranteed to match how its literal was written.
  # Callers building a command (as opposed to an ordinary nested sub-document,
  # where order is irrelevant) must pass a list of {key, value} pairs so the
  # command name is guaranteed to be encoded first.
  defp encode_doc(fields) when is_map(fields) or is_list(fields) do
    body = fields |> Enum.map(fn {k, v} -> encode_elem(to_string(k), v) end) |> :erlang.iolist_to_binary()
    size = 4 + byte_size(body) + 1
    <<size::little-signed-32, body::binary, 0>>
  end

  defp encode_elem(k, v) when is_integer(v) and v >= -2_147_483_648 and v <= 2_147_483_647,
    do: <<0x10, k::binary, 0, v::little-signed-32>>
  defp encode_elem(k, v) when is_integer(v),
    do: <<0x12, k::binary, 0, v::little-signed-64>>
  defp encode_elem(k, v) when is_float(v),
    do: <<0x01, k::binary, 0, v::little-float-64>>
  defp encode_elem(k, v) when is_binary(v) do
    len = byte_size(v) + 1
    <<0x02, k::binary, 0, len::little-32, v::binary, 0>>
  end
  defp encode_elem(k, {:bin, data}) when is_binary(data) do
    <<0x05, k::binary, 0, byte_size(data)::little-32, 0, data::binary>>
  end
  defp encode_elem(k, true),  do: <<0x08, k::binary, 0, 1>>
  defp encode_elem(k, false), do: <<0x08, k::binary, 0, 0>>
  defp encode_elem(k, nil),   do: <<0x0A, k::binary, 0>>
  defp encode_elem(k, v) when is_map(v) do
    doc = encode_doc(v)
    <<0x03, k::binary, 0, doc::binary>>
  end
  defp encode_elem(k, v) when is_list(v) do
    map = v |> Enum.with_index() |> Map.new(fn {val, i} -> {to_string(i), val} end)
    doc = encode_doc(map)
    <<0x04, k::binary, 0, doc::binary>>
  end

  # ---------------------------------------------------------------------------
  # BSON decoder
  # ---------------------------------------------------------------------------

  defp decode_doc(<<size::little-signed-32, rest::binary>>) do
    content_size = size - 5
    <<content::binary-size(^content_size), 0, _::binary>> = rest
    decode_elems(content, %{})
  end
  defp decode_doc(_), do: %{}

  defp decode_elems(<<>>, acc), do: acc
  defp decode_elems(<<type, rest::binary>>, acc) do
    [key, rest2] = :binary.split(rest, <<0>>)
    {value, rest3} = decode_val(type, rest2)
    decode_elems(rest3, Map.put(acc, key, value))
  end

  # double
  defp decode_val(0x01, <<f::little-float-64, r::binary>>), do: {f, r}
  # string
  defp decode_val(0x02, <<len::little-32, r::binary>>) do
    sl = len - 1
    <<s::binary-size(^sl), 0, r2::binary>> = r
    {s, r2}
  end
  # embedded document
  defp decode_val(0x03, <<size::little-signed-32, _::binary>> = data) do
    <<doc::binary-size(^size), r::binary>> = data
    {decode_doc(doc), r}
  end
  # array (same structure as document)
  defp decode_val(0x04, <<size::little-signed-32, _::binary>> = data) do
    <<arr::binary-size(^size), r::binary>> = data
    map  = decode_doc(arr)
    list = map |> Enum.sort_by(fn {k, _} ->
      case Integer.parse(k) do
        {n, _} -> n
        :error  -> 0
      end
    end) |> Enum.map(&elem(&1, 1))
    {list, r}
  end
  # binary
  defp decode_val(0x05, <<len::little-32, sub::8, r::binary>>) do
    <<data::binary-size(^len), r2::binary>> = r
    {{:binary, sub, data}, r2}
  end
  # undefined (deprecated) — no value bytes
  defp decode_val(0x06, r), do: {nil, r}
  # ObjectId — 12 bytes
  defp decode_val(0x07, <<_::binary-12, r::binary>>), do: {nil, r}
  # boolean
  defp decode_val(0x08, <<0, r::binary>>), do: {false, r}
  defp decode_val(0x08, <<_, r::binary>>), do: {true, r}
  # UTC datetime — int64
  defp decode_val(0x09, <<_::little-64, r::binary>>), do: {nil, r}
  # null
  defp decode_val(0x0A, r), do: {nil, r}
  # regex — two cstrings (pattern + flags)
  defp decode_val(0x0B, data) do
    [_, r1] = :binary.split(data, <<0>>)
    [_, r2] = :binary.split(r1,   <<0>>)
    {nil, r2}
  end
  # DBPointer (deprecated) — string + 12 bytes ObjectId
  defp decode_val(0x0C, <<len::little-32, r::binary>>) do
    sl = len - 1
    <<_::binary-size(^sl), 0, _::binary-12, r2::binary>> = r
    {nil, r2}
  end
  # JavaScript code — string format
  defp decode_val(0x0D, <<len::little-32, r::binary>>) do
    sl = len - 1
    <<_::binary-size(^sl), 0, r2::binary>> = r
    {nil, r2}
  end
  # Symbol (deprecated) — string format
  defp decode_val(0x0E, <<len::little-32, r::binary>>) do
    sl = len - 1
    <<_::binary-size(^sl), 0, r2::binary>> = r
    {nil, r2}
  end
  # int32
  defp decode_val(0x10, <<n::little-signed-32, r::binary>>), do: {n, r}
  # Timestamp — uint64
  defp decode_val(0x11, <<_::little-64, r::binary>>), do: {nil, r}
  # int64
  defp decode_val(0x12, <<n::little-signed-64, r::binary>>), do: {n, r}
  # Decimal128 — 16 bytes
  defp decode_val(0x13, <<_::binary-16, r::binary>>), do: {nil, r}
  # MaxKey (0x7F) and MinKey (0xFF) — no value bytes
  defp decode_val(0x7F, r), do: {nil, r}
  defp decode_val(0xFF, r), do: {nil, r}
  # Unknown type — return nil and remaining bytes unchanged
  # (should not occur with well-formed MongoDB responses)
  defp decode_val(_type, r), do: {nil, r}

  # ---------------------------------------------------------------------------
  # Metric helpers
  # ---------------------------------------------------------------------------

  defp total_ops_rate(current, prev, elapsed) when is_number(elapsed) and elapsed > 0 do
    keys = ["insert", "query", "update", "delete", "command"]
    cur  = Enum.reduce(keys, 0, fn k, acc -> acc + (int_val(current[k]) || 0) end)
    prv  = Enum.reduce(keys, 0, fn k, acc -> acc + (int_val(prev[k])    || 0) end)
    if cur >= prv, do: Float.round((cur - prv) / elapsed, 2), else: nil
  end
  defp total_ops_rate(_, _, _), do: nil

  defp op_rate(current, prev, key, elapsed) when is_number(elapsed) and elapsed > 0 do
    cur = int_val(current[key]) || 0
    prv = int_val(prev[key])    || 0
    if cur >= prv, do: Float.round((cur - prv) / elapsed, 2), else: nil
  end
  defp op_rate(_, _, _, _), do: nil

  defp elapsed_sec(nil, _now), do: nil
  defp elapsed_sec(prev, now) do
    diff = DateTime.diff(now, prev, :millisecond)
    if diff > 0, do: diff / 1000, else: nil
  end

  defp int_val(n) when is_integer(n), do: n
  defp int_val(f) when is_float(f),   do: round(f)
  defp int_val(_), do: nil

  defp num_val(n) when is_number(n), do: n
  defp num_val(_), do: nil

  # ---------------------------------------------------------------------------
  # Helpers
  # ---------------------------------------------------------------------------

  defp maybe_put(map, _key, nil), do: map
  defp maybe_put(map, key, val),  do: Map.put(map, key, val)

  defp get_param(params, key) when is_atom(key) do
    v = params[key] || params[to_string(key)]
    if is_binary(v) and String.trim(v) == "", do: nil, else: v
  end

  defp parse_int(nil, default), do: default
  defp parse_int(v, _) when is_integer(v), do: v
  defp parse_int(v, default) when is_binary(v) do
    case Integer.parse(v) do
      {n, _} -> n
      :error -> default
    end
  end
  defp parse_int(_, default), do: default
end