wulf-pulse/lib/services/rmm/persistence.ts

/**
 * Read/write helpers for rmm_executions.
 *
 * Status lifecycle: queued (row inserted, runQuickJob not yet returned) →
 * running (job_uid stored, poller is watching) → complete | failed | timeout.
 */

import postgresClient from '@/lib/services/postgres-client';

// Hard timeout for a dispatched RMM job — covers Datto queue latency, agent
// pickup, script runtime, and result postback. LogLift in particular can run
// long on busy machines (event log collection + gzip + B2 upload). Configurable
// via env so we can bump it without a code change.
const EXECUTION_TIMEOUT_MINUTES = (() => {
  const raw = parseInt(process.env.RMM_EXECUTION_TIMEOUT_MINUTES ?? '', 10);
  return Number.isFinite(raw) && raw > 0 ? raw : 10;
})();

export type RmmExecutionStatus = 'queued' | 'running' | 'complete' | 'failed' | 'timeout';
export type RmmTargetType = 'site_anchor' | 'asset_self';
export type RmmTransport = 'overshell_stdout' | 'b2_upload';

export interface RmmExecutionRow {
  id: string;
  scriptId: string;
  scriptVersion: number;
  targetType: RmmTargetType;
  targetDeviceUid: string;
  targetHostname: string | null;
  targetCompanyId: string | null;
  triggeredByAuditId: string | null;
  assetType: 'flexible_asset' | 'configuration' | null;
  assetId: string | null;
  jobUid: string | null;
  jobName: string;
  variables: unknown;
  status: RmmExecutionStatus;
  exitCode: number | null;
  rawStdout: string | null;
  rawStderr: string | null;
  parsedEvidence: unknown;
  parseError: string | null;
  errorMessage: string | null;
  performedByUserId: string | null;
  transport: RmmTransport;
  evidenceObjectKey: string | null;
  runId: string | null;
  queuedAt: string;
  startedAt: string | null;
  completedAt: string | null;
  timeoutAt: string;
}

interface RawRow {
  id: string;
  script_id: string;
  script_version: number;
  target_type: RmmTargetType;
  target_device_uid: string;
  target_hostname: string | null;
  target_company_id: string | null;
  triggered_by_audit_id: string | null;
  asset_type: 'flexible_asset' | 'configuration' | null;
  asset_id: string | null;
  job_uid: string | null;
  job_name: string;
  variables: unknown;
  status: RmmExecutionStatus;
  exit_code: number | null;
  raw_stdout: string | null;
  raw_stderr: string | null;
  parsed_evidence: unknown;
  parse_error: string | null;
  error_message: string | null;
  performed_by_user_id: string | null;
  transport: RmmTransport;
  evidence_object_key: string | null;
  run_id: string | null;
  queued_at: Date;
  started_at: Date | null;
  completed_at: Date | null;
  timeout_at: Date;
}

const SELECT = `
  id::text                AS id,
  script_id, script_version,
  target_type, target_device_uid, target_hostname,
  target_company_id::text AS target_company_id,
  triggered_by_audit_id::text AS triggered_by_audit_id,
  asset_type, asset_id::text AS asset_id,
  job_uid, job_name, variables,
  status, exit_code, raw_stdout, raw_stderr,
  parsed_evidence, parse_error, error_message,
  performed_by_user_id,
  transport, evidence_object_key, run_id,
  queued_at, started_at, completed_at, timeout_at
`;

function toRow(r: RawRow): RmmExecutionRow {
  return {
    id: r.id,
    scriptId: r.script_id,
    scriptVersion: r.script_version,
    targetType: r.target_type,
    targetDeviceUid: r.target_device_uid,
    targetHostname: r.target_hostname,
    targetCompanyId: r.target_company_id,
    triggeredByAuditId: r.triggered_by_audit_id,
    assetType: r.asset_type,
    assetId: r.asset_id,
    jobUid: r.job_uid,
    jobName: r.job_name,
    variables: r.variables,
    status: r.status,
    exitCode: r.exit_code,
    rawStdout: r.raw_stdout,
    rawStderr: r.raw_stderr,
    parsedEvidence: r.parsed_evidence,
    parseError: r.parse_error,
    errorMessage: r.error_message,
    performedByUserId: r.performed_by_user_id,
    transport: r.transport,
    evidenceObjectKey: r.evidence_object_key,
    runId: r.run_id,
    queuedAt: r.queued_at.toISOString(),
    startedAt: r.started_at?.toISOString() ?? null,
    completedAt: r.completed_at?.toISOString() ?? null,
    timeoutAt: r.timeout_at.toISOString(),
  };
}

export interface CreateExecutionInput {
  scriptId: string;
  scriptVersion: number;
  targetType: RmmTargetType;
  targetDeviceUid: string;
  targetHostname: string | null;
  targetCompanyId: number | string | null;
  triggeredByAuditId?: string | null;
  assetType?: 'flexible_asset' | 'configuration' | null;
  assetId?: number | string | null;
  jobName: string;
  variables: unknown;
  performedByUserId: string | null;
  transport?: RmmTransport;
  runId?: string | null;
}

export async function createPendingExecution(
  input: CreateExecutionInput
): Promise<{ id: string }> {
  const res = await postgresClient.query<{ id: string }>(
    `INSERT INTO rmm_executions
       (script_id, script_version, target_type, target_device_uid,
        target_hostname, target_company_id,
        triggered_by_audit_id, asset_type, asset_id,
        job_name, variables, status, performed_by_user_id,
        transport, run_id, timeout_at)
     VALUES ($1, $2, $3, $4,
             $5, $6,
             $7, $8, $9,
             $10, $11::jsonb, 'queued', $12,
             $13, $14, NOW() + ($15::int || ' minutes')::interval)
     RETURNING id::text AS id`,
    [
      input.scriptId,
      input.scriptVersion,
      input.targetType,
      input.targetDeviceUid,
      input.targetHostname,
      input.targetCompanyId,
      input.triggeredByAuditId ?? null,
      input.assetType ?? null,
      input.assetId ?? null,
      input.jobName,
      JSON.stringify(input.variables ?? {}),
      input.performedByUserId,
      input.transport ?? 'overshell_stdout',
      input.runId ?? null,
      EXECUTION_TIMEOUT_MINUTES,
    ]
  );
  return { id: res.rows[0].id };
}

export async function markExecutionRunning(
  id: string,
  jobUid: string
): Promise<void> {
  await postgresClient.query(
    `UPDATE rmm_executions
        SET status = 'running',
            job_uid = $2,
            started_at = NOW()
      WHERE id = $1`,
    [id, jobUid]
  );
}

export async function markExecutionFailedToDispatch(
  id: string,
  message: string
): Promise<void> {
  await postgresClient.query(
    `UPDATE rmm_executions
        SET status = 'failed',
            error_message = $2,
            completed_at = NOW()
      WHERE id = $1`,
    [id, message]
  );
}

export async function markExecutionComplete(input: {
  id: string;
  exitCode: number | null;
  rawStdout: string | null;
  rawStderr: string | null;
  parsedEvidence: unknown;
  parseError: string | null;
}): Promise<void> {
  const finalStatus = input.exitCode !== null && input.exitCode !== 0 ? 'failed' : 'complete';
  await postgresClient.query(
    `UPDATE rmm_executions
        SET status = $2,
            exit_code = $3,
            raw_stdout = $4,
            raw_stderr = $5,
            parsed_evidence = $6::jsonb,
            parse_error = $7,
            completed_at = NOW()
      WHERE id = $1`,
    [
      input.id,
      finalStatus,
      input.exitCode,
      input.rawStdout,
      input.rawStderr,
      input.parsedEvidence === undefined ? null : JSON.stringify(input.parsedEvidence ?? null),
      input.parseError,
    ]
  );
}

export async function markExecutionTimeout(id: string): Promise<void> {
  await postgresClient.query(
    `UPDATE rmm_executions
        SET status = 'timeout',
            completed_at = NOW(),
            error_message = COALESCE(error_message,
              'Hard timeout reached — Datto RMM did not return job results within ' || $2::text || ' minutes.')
      WHERE id = $1`,
    [id, EXECUTION_TIMEOUT_MINUTES]
  );
}

export async function getExecutionById(id: string): Promise<RmmExecutionRow | null> {
  const res = await postgresClient.query<RawRow>(
    `SELECT ${SELECT} FROM rmm_executions WHERE id = $1 LIMIT 1`,
    [id]
  );
  if (res.rowCount === 0) return null;
  return toRow(res.rows[0]);
}

/** Pull rows the worker should poll (running, not yet timed out). */
export async function listRunningExecutions(): Promise<RmmExecutionRow[]> {
  const res = await postgresClient.query<RawRow>(
    `SELECT ${SELECT} FROM rmm_executions
      WHERE status = 'running' AND timeout_at > NOW()
      ORDER BY queued_at`
  );
  return res.rows.map(toRow);
}

/** Pull rows that have passed their timeout marker. */
export async function listTimedOutExecutions(): Promise<RmmExecutionRow[]> {
  const res = await postgresClient.query<RawRow>(
    `SELECT ${SELECT} FROM rmm_executions
      WHERE status IN ('queued','running')
        AND timeout_at <= NOW()
      ORDER BY queued_at`
  );
  return res.rows.map(toRow);
}

/** List recent executions. Filters: company, asset, script, status. */
export async function listExecutions(opts: {
  limit?: number;
  offset?: number;
  companyId?: number | string;
  scriptId?: string;
  status?: RmmExecutionStatus;
  assetType?: 'flexible_asset' | 'configuration';
  assetId?: number | string;
}): Promise<RmmExecutionRow[]> {
  const limit = Math.min(opts.limit ?? 100, 500);
  const offset = opts.offset ?? 0;
  const params: unknown[] = [limit, offset];
  const where: string[] = [];
  if (opts.companyId !== undefined) {
    params.push(opts.companyId);
    where.push(`target_company_id = $${params.length}`);
  }
  if (opts.scriptId) {
    params.push(opts.scriptId);
    where.push(`script_id = $${params.length}`);
  }
  if (opts.status) {
    params.push(opts.status);
    where.push(`status = $${params.length}`);
  }
  if (opts.assetType) {
    params.push(opts.assetType);
    where.push(`asset_type = $${params.length}`);
  }
  if (opts.assetId !== undefined) {
    params.push(opts.assetId);
    where.push(`asset_id = $${params.length}`);
  }
  const whereSql = where.length === 0 ? '' : `WHERE ${where.join(' AND ')}`;
  const res = await postgresClient.query<RawRow>(
    `SELECT ${SELECT}
       FROM rmm_executions
       ${whereSql}
      ORDER BY queued_at DESC
      LIMIT $1 OFFSET $2`,
    params
  );
  return res.rows.map(toRow);
}

/**
 * For the audit pipeline: for a given company, return the latest successful
 * execution per script_id within the last `since` days.
 */
export async function listLatestEvidenceForCompany(
  companyId: number | string,
  sinceDays = 7
): Promise<RmmExecutionRow[]> {
  const res = await postgresClient.query<RawRow>(
    `SELECT DISTINCT ON (script_id) ${SELECT}
       FROM rmm_executions
      WHERE target_company_id = $1
        AND status = 'complete'
        AND completed_at >= NOW() - ($2::int || ' days')::interval
      ORDER BY script_id, completed_at DESC`,
    [companyId, sinceDays]
  );
  return res.rows.map(toRow);
}

/**
 * For the audit pipeline: for a specific asset, return the latest successful
 * asset-self execution per script_id (no time limit — asset-self evidence
 * stays relevant longer than site-anchored).
 */
export async function listLatestEvidenceForAsset(
  assetType: 'flexible_asset' | 'configuration',
  assetId: number | string
): Promise<RmmExecutionRow[]> {
  const res = await postgresClient.query<RawRow>(
    `SELECT DISTINCT ON (script_id) ${SELECT}
       FROM rmm_executions
      WHERE asset_type = $1
        AND asset_id = $2
        AND target_type = 'asset_self'
        AND status = 'complete'
      ORDER BY script_id, completed_at DESC`,
    [assetType, assetId]
  );
  return res.rows.map(toRow);
}

/**
 * Webhook correlation: find a still-pending b2_upload execution dispatched
 * by Pulse with a given run_id. Caller updates this row instead of inserting
 * a fresh one when an inbound LogLift upload matches a Pulse-driven dispatch.
 */
export async function findExecutionByRunId(
  runId: string
): Promise<RmmExecutionRow | null> {
  const res = await postgresClient.query<RawRow>(
    `SELECT ${SELECT}
       FROM rmm_executions
      WHERE run_id = $1
      ORDER BY queued_at DESC
      LIMIT 1`,
    [runId]
  );
  if (res.rowCount === 0) return null;
  return toRow(res.rows[0]);
}

/**
 * Insert a fresh row for an inbound LogLift webhook that didn't correlate to a
 * Pulse-dispatched execution (e.g. the collector ran on its own schedule).
 * Status starts at 'running' — the webhook is the completion event.
 */
export async function createOutOfBandUploadExecution(input: {
  scriptId: string;
  scriptVersion: number;
  runId: string;
  jobName: string;
  targetDeviceUid: string;
  targetHostname: string | null;
  targetCompanyId: number | string | null;
  assetType: 'flexible_asset' | 'configuration' | null;
  assetId: number | string | null;
  variables: unknown;
}): Promise<{ id: string }> {
  const res = await postgresClient.query<{ id: string }>(
    `INSERT INTO rmm_executions
       (script_id, script_version, target_type, target_device_uid,
        target_hostname, target_company_id,
        asset_type, asset_id,
        job_name, variables, status, performed_by_user_id,
        transport, run_id, started_at)
     VALUES ($1, $2, 'asset_self', $3,
             $4, $5,
             $6, $7,
             $8, $9::jsonb, 'running', NULL,
             'b2_upload', $10, NOW())
     RETURNING id::text AS id`,
    [
      input.scriptId,
      input.scriptVersion,
      input.targetDeviceUid,
      input.targetHostname,
      input.targetCompanyId,
      input.assetType,
      input.assetId,
      input.jobName,
      JSON.stringify(input.variables ?? {}),
      input.runId,
    ]
  );
  return { id: res.rows[0].id };
}

/**
 * Mark a b2_upload execution complete with parsed evidence + the B2 object key
 * for forensic replay. Caller is responsible for resolving target_company_id
 * and asset linkage before calling (the webhook receiver does this from the
 * payload's clientId / computerName).
 */
export async function markExecutionFromB2Upload(input: {
  id: string;
  evidenceObjectKey: string;
  parsedEvidence: unknown;
  targetCompanyId?: number | string | null;
  assetType?: 'flexible_asset' | 'configuration' | null;
  assetId?: number | string | null;
}): Promise<void> {
  await postgresClient.query(
    `UPDATE rmm_executions
        SET status = 'complete',
            exit_code = 0,
            evidence_object_key = $2,
            parsed_evidence = $3::jsonb,
            target_company_id = COALESCE($4, target_company_id),
            asset_type = COALESCE($5, asset_type),
            asset_id = COALESCE($6, asset_id),
            completed_at = NOW()
      WHERE id = $1`,
    [
      input.id,
      input.evidenceObjectKey,
      JSON.stringify(input.parsedEvidence ?? null),
      input.targetCompanyId ?? null,
      input.assetType ?? null,
      input.assetId ?? null,
    ]
  );
}

/** Per-user 24-hour count for rate-limiting. */
export async function countUserExecutionsLast24h(userId: string): Promise<number> {
  const res = await postgresClient.query<{ count: string }>(
    `SELECT COUNT(*)::text AS count
       FROM rmm_executions
      WHERE performed_by_user_id = $1
        AND queued_at >= NOW() - INTERVAL '24 hours'`,
    [userId]
  );
  return Number(res.rows[0]?.count ?? 0);
}