simstudioai--sim/apps/sim/app/api/tools/postgresql/utils.ts

import postgres from 'postgres'
import { validateDatabaseHost } from '@/lib/core/security/input-validation.server'
import type { PostgresConnectionConfig } from '@/tools/postgresql/types'

export async function createPostgresConnection(config: PostgresConnectionConfig) {
  const hostValidation = await validateDatabaseHost(config.host, 'host')
  if (!hostValidation.isValid) {
    throw new Error(hostValidation.error)
  }

  const resolvedHost = hostValidation.resolvedIP ?? config.host

  const sslConfig: boolean | 'prefer' | { rejectUnauthorized: boolean; servername?: string } =
    config.ssl === 'disabled'
      ? false
      : config.ssl === 'preferred'
        ? 'prefer'
        : { rejectUnauthorized: false, servername: config.host }

  const sql = postgres({
    // Pin the validated IP (never the hostname) to prevent DNS rebinding; SNI stays the hostname above.
    host: resolvedHost,
    port: config.port,
    database: config.database,
    username: config.username,
    password: config.password,
    ssl: sslConfig,
    connect_timeout: 10, // 10 seconds
    idle_timeout: 20, // 20 seconds
    max_lifetime: 60 * 30, // 30 minutes
    max: 1, // Single connection for tool usage
  })

  return sql
}

export async function executeQuery(
  sql: any,
  query: string,
  params: unknown[] = []
): Promise<{ rows: unknown[]; rowCount: number }> {
  const result = await sql.unsafe(query, params)
  const rowCount = result.count ?? result.length ?? 0
  return {
    rows: Array.isArray(result) ? result : [result],
    rowCount,
  }
}

export function validateQuery(query: string): { isValid: boolean; error?: string } {
  const trimmedQuery = query.trim().toLowerCase()

  const allowedStatements = /^(select|insert|update|delete|with|explain|analyze|show)\s+/i
  if (!allowedStatements.test(trimmedQuery)) {
    return {
      isValid: false,
      error:
        'Only SELECT, INSERT, UPDATE, DELETE, WITH, EXPLAIN, ANALYZE, and SHOW statements are allowed',
    }
  }

  return { isValid: true }
}

export function sanitizeIdentifier(identifier: string): string {
  if (identifier.includes('.')) {
    const parts = identifier.split('.')
    return parts.map((part) => sanitizeSingleIdentifier(part)).join('.')
  }

  return sanitizeSingleIdentifier(identifier)
}

/**
 * Validates a WHERE clause to prevent SQL injection attacks
 * @param where - The WHERE clause string to validate
 * @throws {Error} If the WHERE clause contains potentially dangerous patterns
 */
function validateWhereClause(where: string): void {
  const dangerousPatterns = [
    // DDL and DML injection via stacked queries
    /;\s*(drop|delete|insert|update|create|alter|grant|revoke)/i,
    // Union-based injection
    /union\s+(all\s+)?select/i,
    // File operations
    /into\s+outfile/i,
    /load_file\s*\(/i,
    /pg_read_file/i,
    // Comment-based injection (can truncate query)
    /--/,
    /\/\*/,
    /\*\//,
    // Tautologies - always true/false conditions using backreferences
    // Matches OR 'x'='x' or OR x=x (same value both sides) but NOT OR col='value'
    /\bor\s+(['"]?)(\w+)\1\s*=\s*\1\2\1/i,
    /\bor\s+true\b/i,
    /\bor\s+false\b/i,
    // AND tautologies (less common but still used in attacks)
    /\band\s+(['"]?)(\w+)\1\s*=\s*\1\2\1/i,
    /\band\s+true\b/i,
    /\band\s+false\b/i,
    // Time-based blind injection
    /\bsleep\s*\(/i,
    /\bwaitfor\s+delay/i,
    /\bpg_sleep\s*\(/i,
    /\bbenchmark\s*\(/i,
    // Stacked queries (any statement after semicolon)
    /;\s*\w+/,
    // Information schema / system catalog queries
    /information_schema/i,
    /pg_catalog/i,
    // System functions and procedures
    /\bxp_cmdshell/i,
  ]

  for (const pattern of dangerousPatterns) {
    if (pattern.test(where)) {
      throw new Error('WHERE clause contains potentially dangerous operation')
    }
  }
}

function sanitizeSingleIdentifier(identifier: string): string {
  const cleaned = identifier.replace(/"/g, '')

  if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(cleaned)) {
    throw new Error(
      `Invalid identifier: ${identifier}. Identifiers must start with a letter or underscore and contain only letters, numbers, and underscores.`
    )
  }

  return `"${cleaned}"`
}

export async function executeInsert(
  sql: any,
  table: string,
  data: Record<string, unknown>
): Promise<{ rows: unknown[]; rowCount: number }> {
  const sanitizedTable = sanitizeIdentifier(table)
  const columns = Object.keys(data)
  const sanitizedColumns = columns.map((col) => sanitizeIdentifier(col))
  const placeholders = columns.map((_, index) => `$${index + 1}`)
  const values = columns.map((col) => data[col])

  const query = `INSERT INTO ${sanitizedTable} (${sanitizedColumns.join(', ')}) VALUES (${placeholders.join(', ')}) RETURNING *`
  const result = await sql.unsafe(query, values)

  const rowCount = result.count ?? result.length ?? 0
  return {
    rows: Array.isArray(result) ? result : [result],
    rowCount,
  }
}

export async function executeUpdate(
  sql: any,
  table: string,
  data: Record<string, unknown>,
  where: string
): Promise<{ rows: unknown[]; rowCount: number }> {
  validateWhereClause(where)

  const sanitizedTable = sanitizeIdentifier(table)
  const columns = Object.keys(data)
  const sanitizedColumns = columns.map((col) => sanitizeIdentifier(col))
  const setClause = sanitizedColumns.map((col, index) => `${col} = $${index + 1}`).join(', ')
  const values = columns.map((col) => data[col])

  const query = `UPDATE ${sanitizedTable} SET ${setClause} WHERE ${where} RETURNING *`
  const result = await sql.unsafe(query, values)

  const rowCount = result.count ?? result.length ?? 0
  return {
    rows: Array.isArray(result) ? result : [result],
    rowCount,
  }
}

export async function executeDelete(
  sql: any,
  table: string,
  where: string
): Promise<{ rows: unknown[]; rowCount: number }> {
  validateWhereClause(where)

  const sanitizedTable = sanitizeIdentifier(table)
  const query = `DELETE FROM ${sanitizedTable} WHERE ${where} RETURNING *`
  const result = await sql.unsafe(query, [])

  const rowCount = result.count ?? result.length ?? 0
  return {
    rows: Array.isArray(result) ? result : [result],
    rowCount,
  }
}

export interface IntrospectionResult {
  tables: Array<{
    name: string
    schema: string
    columns: Array<{
      name: string
      type: string
      nullable: boolean
      default: string | null
      isPrimaryKey: boolean
      isForeignKey: boolean
      references?: {
        table: string
        column: string
      }
    }>
    primaryKey: string[]
    foreignKeys: Array<{
      column: string
      referencesTable: string
      referencesColumn: string
    }>
    indexes: Array<{
      name: string
      columns: string[]
      unique: boolean
    }>
  }>
  schemas: string[]
}

export async function executeIntrospect(
  sql: any,
  schemaName = 'public'
): Promise<IntrospectionResult> {
  const schemasResult = await sql`
    SELECT schema_name
    FROM information_schema.schemata
    WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
    ORDER BY schema_name
  `
  const schemas = schemasResult.map((row: { schema_name: string }) => row.schema_name)

  const tablesResult = await sql`
    SELECT table_name, table_schema
    FROM information_schema.tables
    WHERE table_schema = ${schemaName}
      AND table_type = 'BASE TABLE'
    ORDER BY table_name
  `

  const tables = []

  for (const tableRow of tablesResult) {
    const tableName = tableRow.table_name
    const tableSchema = tableRow.table_schema

    const columnsResult = await sql`
      SELECT
        c.column_name,
        c.data_type,
        c.is_nullable,
        c.column_default,
        c.udt_name
      FROM information_schema.columns c
      WHERE c.table_schema = ${tableSchema}
        AND c.table_name = ${tableName}
      ORDER BY c.ordinal_position
    `

    const pkResult = await sql`
      SELECT kcu.column_name
      FROM information_schema.table_constraints tc
      JOIN information_schema.key_column_usage kcu
        ON tc.constraint_name = kcu.constraint_name
        AND tc.table_schema = kcu.table_schema
      WHERE tc.constraint_type = 'PRIMARY KEY'
        AND tc.table_schema = ${tableSchema}
        AND tc.table_name = ${tableName}
    `
    const primaryKeyColumns = pkResult.map((row: { column_name: string }) => row.column_name)

    const fkResult = await sql`
      SELECT
        kcu.column_name,
        ccu.table_name AS foreign_table_name,
        ccu.column_name AS foreign_column_name
      FROM information_schema.table_constraints tc
      JOIN information_schema.key_column_usage kcu
        ON tc.constraint_name = kcu.constraint_name
        AND tc.table_schema = kcu.table_schema
      JOIN information_schema.constraint_column_usage ccu
        ON ccu.constraint_name = tc.constraint_name
        AND ccu.table_schema = tc.table_schema
      WHERE tc.constraint_type = 'FOREIGN KEY'
        AND tc.table_schema = ${tableSchema}
        AND tc.table_name = ${tableName}
    `

    const foreignKeys = fkResult.map(
      (row: { column_name: string; foreign_table_name: string; foreign_column_name: string }) => ({
        column: row.column_name,
        referencesTable: row.foreign_table_name,
        referencesColumn: row.foreign_column_name,
      })
    )

    const fkColumnSet = new Set(foreignKeys.map((fk: { column: string }) => fk.column))

    const indexesResult = await sql`
      SELECT
        i.relname AS index_name,
        a.attname AS column_name,
        ix.indisunique AS is_unique
      FROM pg_class t
      JOIN pg_index ix ON t.oid = ix.indrelid
      JOIN pg_class i ON i.oid = ix.indexrelid
      JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
      JOIN pg_namespace n ON n.oid = t.relnamespace
      WHERE t.relkind = 'r'
        AND n.nspname = ${tableSchema}
        AND t.relname = ${tableName}
        AND NOT ix.indisprimary
      ORDER BY i.relname, a.attnum
    `

    const indexMap = new Map<string, { name: string; columns: string[]; unique: boolean }>()
    for (const row of indexesResult) {
      const indexName = row.index_name
      if (!indexMap.has(indexName)) {
        indexMap.set(indexName, {
          name: indexName,
          columns: [],
          unique: row.is_unique,
        })
      }
      indexMap.get(indexName)!.columns.push(row.column_name)
    }
    const indexes = Array.from(indexMap.values())

    const columns = columnsResult.map(
      (col: {
        column_name: string
        data_type: string
        is_nullable: string
        column_default: string | null
        udt_name: string
      }) => {
        const columnName = col.column_name
        const fk = foreignKeys.find((f: { column: string }) => f.column === columnName)

        return {
          name: columnName,
          type: col.data_type === 'USER-DEFINED' ? col.udt_name : col.data_type,
          nullable: col.is_nullable === 'YES',
          default: col.column_default,
          isPrimaryKey: primaryKeyColumns.includes(columnName),
          isForeignKey: fkColumnSet.has(columnName),
          ...(fk && {
            references: {
              table: fk.referencesTable,
              column: fk.referencesColumn,
            },
          }),
        }
      }
    )

    tables.push({
      name: tableName,
      schema: tableSchema,
      columns,
      primaryKey: primaryKeyColumns,
      foreignKeys,
      indexes,
    })
  }

  return { tables, schemas }
}