feistel_cipher

defmodule FeistelCipher do
  @moduledoc """
  Encrypted integer IDs using Feistel cipher.

  ## Basic Usage

  ```elixir
  defmodule MyApp.Repo.Migrations.AddFeistelCipher do
    use Ecto.Migration

    def up, do: FeistelCipher.up_v1_functions()
    def down, do: FeistelCipher.down_v1_functions()
  end
  ```

  ## With Custom Prefix

  ```elixir
  def up, do: FeistelCipher.up_v1_functions(functions_prefix: "private")
  def down, do: FeistelCipher.down_v1_functions(functions_prefix: "private")
  ```

  See function documentation for detailed options and examples.
  """

  use Ecto.Migration

  @backfill_sentinel -1

  @doc """
  Generates a cryptographically secure random salt for the Feistel cipher.

  Returns a random integer between 0 and 2^31-1.

  ## Example

      salt = FeistelCipher.generate_random_salt()
      FeistelCipher.up_v1_functions(functions_salt: salt)
  """
  @spec generate_random_salt() :: non_neg_integer()
  def generate_random_salt do
    <<salt::31, _::1>> = :crypto.strong_rand_bytes(4)
    salt
  end

  @doc """
  Create FeistelCipher functions in the database.

  ## Options

  * `:functions_prefix` - Schema prefix for functions (default: "public"). Can be changed, but should be treated as an explicit migration because trigger SQL must reference the same function schema.
  * `:functions_salt` - Salt constant for cipher algorithm. A random value is generated by default if not specified. Must be 0 to 2^31-1. Can be changed, but should be treated as an explicit migration because cipher output compatibility changes.
  """
  @spec up_v1_functions(keyword()) :: :ok
  def up_v1_functions(opts \\ []) when is_list(opts) do
    functions_prefix = Keyword.get(opts, :functions_prefix, "public")
    functions_salt = Keyword.get(opts, :functions_salt, generate_random_salt())
    validate_key!(functions_salt, "functions_salt")

    execute("CREATE SCHEMA IF NOT EXISTS #{functions_prefix}")
    execute("CREATE EXTENSION IF NOT EXISTS pgcrypto")

    # Feistel cipher implementation based on https://wiki.postgresql.org/wiki/Pseudo_encrypt
    # Enhanced with key-based HMAC-SHA256 for cryptographic security
    # Algorithm reference: https://www.youtube.com/watch?v=FGhj3CGxl8I

    # bigint is 64 bits, but excluding negative numbers, only 63 bits are usable.
    # Limited to a maximum of 62 bits as it needs to be halved for the operation.
    # Multiplication and operation parameters are all limited to 31 bits.
    # Since 31 bits (half of 62 bits) are multiplied by a 31-bit parameter,
    # the calculation result is also within the 62-bit range, making it safe for bigint.
    execute("""
    CREATE OR REPLACE FUNCTION #{functions_prefix}.feistel_cipher_v1(value bigint, bits int, key bigint, rounds int) returns bigint AS $$
      DECLARE
        round      int;
        left_half  bigint;
        right_half bigint;
        temp       bigint;

        half_bits  int    := bits / 2;
        half_mask  bigint := (1::bigint << half_bits) - 1;
        mask       bigint := (1::bigint << bits) - 1;
        hash_bytes bytea;
        hash_int   bigint;

      BEGIN
        IF bits < 0 OR bits > 62 OR bits % 2 = 1 THEN
          RAISE EXCEPTION 'feistel_cipher_v1: bits must be an even number between 0 and 62: %', bits;
        END IF;

        IF key < 0 OR key >= (1::bigint << 31) THEN
          RAISE EXCEPTION 'feistel_cipher_v1: key must be between 0 and 2^31-1: %', key;
        END IF;

        IF value < 0 OR value > mask THEN
          RAISE EXCEPTION 'feistel_cipher_v1: value must be between 0 and % for % bits: %', mask, bits, value;
        END IF;

        IF rounds < 1 OR rounds > 32 THEN
          RAISE EXCEPTION 'feistel_cipher_v1: rounds must be between 1 and 32: %', rounds;
        END IF;

        -- Split value into left and right halves
        left_half  := (value >> half_bits) & half_mask;
        right_half := value & half_mask;

        -- Feistel rounds with HMAC-SHA256 round function
        -- Using cryptographic HMAC makes the cipher resistant to all known attacks
        -- Note: Round number is NOT included in hash to maintain encrypt/decrypt symmetry
        FOR round IN 1..rounds LOOP
          temp       := right_half;
          hash_bytes := hmac(int4send(right_half::int4), int4send(key::int4) || int4send(#{functions_salt}::int4), 'sha256');
          hash_int   := get_byte(hash_bytes, 0)::bigint << 24
                      | get_byte(hash_bytes, 1)::bigint << 16
                      | get_byte(hash_bytes, 2)::bigint << 8
                      | get_byte(hash_bytes, 3)::bigint;
          right_half := left_half # (hash_int & half_mask);
          left_half  := temp;
        END LOOP;

        -- Final swap
        temp       := left_half;
        left_half  := right_half;
        right_half := temp;

        -- Combine halves
        RETURN ((left_half << half_bits) | right_half);
      END;
    $$ LANGUAGE plpgsql strict immutable;
    """)

    execute("""
    CREATE OR REPLACE FUNCTION #{functions_prefix}.feistel_trigger_v1() RETURNS trigger AS $$
      DECLARE
        -- Trigger parameters
        from_column  text;
        to_column    text;
        time_bits    int;
        time_bucket  bigint;
        time_offset  bigint;
        encrypt_time boolean;
        data_bits    int;
        key          bigint;
        rounds       int;

        -- From and to values
        from_value     bigint;
        to_value       bigint;

        -- Time-related values
        time_value     bigint;
        time_component bigint;

        -- Data-related values
        data_component bigint;

        -- Temporary values for validation
        decrypted      bigint;
        old_from_value bigint;

      BEGIN
        -- Extract trigger parameters
        from_column  := TG_ARGV[0]::text;
        to_column    := TG_ARGV[1]::text;
        time_bits    := TG_ARGV[2]::int;
        time_bucket  := TG_ARGV[3]::bigint;
        time_offset  := TG_ARGV[4]::bigint;
        encrypt_time := TG_ARGV[5]::boolean;
        data_bits    := TG_ARGV[6]::int;
        key          := TG_ARGV[7]::bigint;
        rounds       := TG_ARGV[8]::int;

        -- Early return: If from_column is not modified on UPDATE, skip re-encryption.
        -- This allows manual modification of to_column if from_column remains unchanged.
        IF TG_OP = 'UPDATE' THEN
          EXECUTE format('SELECT ($1).%I::bigint, ($2).%I::bigint', from_column, from_column)
          INTO old_from_value, from_value
          USING OLD, NEW;

          IF old_from_value IS NOT DISTINCT FROM from_value THEN
            RETURN NEW;
          END IF;
        ELSE
          -- Get the from_value from the from_column (for INSERT)
          EXECUTE format('SELECT ($1).%I::bigint', from_column)
          INTO from_value
          USING NEW;
        END IF;

        -- Handle NULL case early
        IF from_value IS NULL THEN
          to_value := NULL;
        ELSE
          -- Calculate time component
          -- When time_bits = 0, all time-related options are effectively ignored.
          IF time_bits = 0 THEN
            time_component := 0;
          ELSE
            time_value :=
              floor((extract(epoch from now()) + time_offset::double precision) / time_bucket::double precision)::bigint;
            time_value := time_value & ((1::bigint << time_bits) - 1);

            IF encrypt_time THEN
              time_component := #{functions_prefix}.feistel_cipher_v1(time_value, time_bits, key, rounds);
            ELSE
              time_component := time_value;
            END IF;
          END IF;

          -- Encrypt the data part using feistel cipher
          data_component := #{functions_prefix}.feistel_cipher_v1(from_value, data_bits, key, rounds);

          -- Sanity check: Verify data part encryption is reversible
          decrypted := #{functions_prefix}.feistel_cipher_v1(data_component, data_bits, key, rounds);

          IF decrypted IS DISTINCT FROM from_value THEN
            RAISE EXCEPTION 'feistel_trigger_v1: feistel_cipher_v1 does not have an inverse (from: %, data_component: %, decrypted: %, data_bits: %, key: %, rounds: %)',
              from_value, data_component, decrypted, data_bits, key, rounds;
          END IF;

          -- Combine: [time_bits | data_bits]
          to_value := (time_component << data_bits) | data_component;
        END IF;

        -- Set the to_value to the to_column in the NEW record
        NEW := jsonb_populate_record(NEW, jsonb_build_object(to_column, to_jsonb(to_value)));

        RETURN NEW;
      END;
    $$ LANGUAGE plpgsql;
    """)
  end

  @doc """
  Drop FeistelCipher functions from the database.

  **Note**: PostgreSQL will automatically prevent this operation if any triggers
  are still using these functions. Drop all triggers first using `down_for_trigger/4`.

  ## Options

  * `:functions_prefix` - Schema prefix where functions are located (default: "public").
  """
  @spec down_v1_functions(keyword()) :: :ok
  def down_v1_functions(opts \\ []) when is_list(opts) do
    functions_prefix = Keyword.get(opts, :functions_prefix, "public")

    execute(
      "DROP FUNCTION IF EXISTS #{functions_prefix}.feistel_cipher_v1(bigint, int, bigint, int)"
    )

    execute("DROP FUNCTION IF EXISTS #{functions_prefix}.feistel_trigger_v1()")
  end

  @doc """
  Creates legacy v0.x PostgreSQL functions (`feistel` and `handle_feistel_encryption`).

  Use this in old migrations that previously called `FeistelCipher.Migration.up/1`
  so that `mix ecto.migrate` works from scratch with feistel_cipher v1.0.

  The created legacy functions are compatibility stubs that raise if called.
  They exist only so historical migrations can run.

  Uses `CREATE OR REPLACE` for idempotency.

  ## Options

  * `:functions_prefix` - Schema prefix (default: "public").
  * `:functions_salt` - Required for backward compatibility with old migrations (unused by stubs).
  """
  @spec up_legacy_functions(keyword()) :: :ok
  def up_legacy_functions(opts \\ []) when is_list(opts) do
    functions_prefix = Keyword.get(opts, :functions_prefix, "public")
    _functions_salt = Keyword.fetch!(opts, :functions_salt)

    execute("CREATE EXTENSION IF NOT EXISTS pgcrypto")

    execute("""
    CREATE OR REPLACE FUNCTION #{functions_prefix}.feistel(input bigint, bits int, key bigint) returns bigint AS $$
      BEGIN
        RAISE EXCEPTION
          'legacy function %.feistel(bigint, int, bigint) is compatibility-only and must not be called; migrate triggers to feistel_trigger_v1',
          #{inspect(functions_prefix)};
      END;
    $$ LANGUAGE plpgsql strict immutable;
    """)

    execute("""
    CREATE OR REPLACE FUNCTION #{functions_prefix}.handle_feistel_encryption() RETURNS trigger AS $$
      BEGIN
        RAISE EXCEPTION
          'legacy trigger function %.handle_feistel_encryption() is compatibility-only and must not be called; migrate triggers to feistel_trigger_v1',
          #{inspect(functions_prefix)};
        RETURN NEW;
      END;
    $$ LANGUAGE plpgsql;
    """)
  end

  @doc """
  Drops legacy v0.x PostgreSQL functions (`feistel` and `handle_feistel_encryption`).

  Uses `IF EXISTS` for idempotency.

  ## Options

  * `:functions_prefix` - Schema prefix (default: "public").
  """
  @spec down_legacy_functions(keyword()) :: :ok
  def down_legacy_functions(opts \\ []) when is_list(opts) do
    functions_prefix = Keyword.get(opts, :functions_prefix, "public")

    execute("DROP FUNCTION IF EXISTS #{functions_prefix}.handle_feistel_encryption()")
    execute("DROP FUNCTION IF EXISTS #{functions_prefix}.feistel(bigint, int, bigint)")
  end

  @doc """
  Returns SQL to create a trigger that encrypts a `from` column to a `to` column.

  The resulting ID structure is `[time_bits | data_bits]`, where time_bits serve as a
  prefix to cluster rows created in the same time bucket on nearby PostgreSQL pages,
  optimizing incremental backup efficiency.

  ## Options

  * `:time_bits` - Time prefix bits (default: 15). Set to 0 for no time prefix. Can be changed, but should be treated as an explicit migration because old/new IDs will use different time-prefix semantics.
  * `:time_bucket` - Time bucket size in seconds (default: 86400 = 1 day). Can be changed, but should be treated as an explicit migration because clustering behavior changes immediately.
  * `:time_offset` - Time offset in seconds applied before bucket calculation (default: 0). For example, 21600 shifts a daily boundary from 00:00 UTC to 18:00 UTC (03:00 KST). Can be changed, but should be treated as an explicit migration because clustering behavior changes immediately.
  * `:encrypt_time` - Whether to encrypt time_bits with feistel cipher (default: false). When true, time_bits must be even. Ignored by the trigger when `time_bits` is `0`. Can be changed, but should be treated as an explicit migration because time-prefix interpretation changes.
  * `:data_bits` - Data cipher bits (default: 38, must be even). Should be treated as immutable in-place; changing it requires a planned migration.
  * `:key` - Encryption key (0 to 2^31-1). Auto-generated if not provided. Should be treated as immutable in-place; changing it requires a planned migration.
  * `:rounds` - Number of Feistel rounds (default: 16, min: 1, max: 32). Should be treated as immutable in-place; changing it requires a planned migration.
  * `:functions_prefix` - Schema where cipher functions are located (default: "public"). Can be changed, but should be treated as an explicit migration because triggers must call the correct schema.

  ## Examples

      FeistelCipher.up_for_legacy_trigger("public", "posts", "seq", "id")

      FeistelCipher.up_for_legacy_trigger("public", "posts", "seq", "id",
        time_bits: 8, time_bucket: 86400, data_bits: 32)

      FeistelCipher.up_for_legacy_trigger("public", "posts", "seq", "id",
        time_bits: 0)

  """
  @spec up_for_legacy_trigger(String.t(), String.t(), String.t(), String.t(), keyword()) ::
          String.t()
  def up_for_legacy_trigger(prefix, table, from, to, opts \\ []) when is_list(opts) do
    do_up_for_trigger(prefix, table, from, to, opts, &legacy_trigger_name/3)
  end

  @doc """
  Same as `up_for_legacy_trigger/5` but uses v1 trigger naming convention.
  """
  @spec up_for_v1_trigger(String.t(), String.t(), String.t(), String.t(), keyword()) ::
          String.t()
  def up_for_v1_trigger(prefix, table, from, to, opts \\ []) when is_list(opts) do
    do_up_for_trigger(prefix, table, from, to, opts, &trigger_name/3)
  end

  defp do_up_for_trigger(prefix, table, from, to, opts, name_fn) do
    resolved_opts = resolve_cipher_options(prefix, table, from, to, opts)

    """
    CREATE TRIGGER #{name_fn.(table, from, to)}
      BEFORE INSERT OR UPDATE
      ON #{prefix}.#{table}
      FOR EACH ROW
      EXECUTE PROCEDURE #{resolved_opts.functions_prefix}.feistel_trigger_v1('#{from}', '#{to}', #{resolved_opts.time_bits}, #{resolved_opts.time_bucket}, #{resolved_opts.time_offset}, #{resolved_opts.encrypt_time}, #{resolved_opts.data_bits}, #{resolved_opts.key}, #{resolved_opts.rounds});
    """
  end

  @doc """
  Returns SQL to backfill a v1 encrypted column for existing rows.

  This is useful when adding a new encrypted column to a table that already has data.
  New rows will be handled by the trigger, and this statement fills only rows where
  the encrypted target column is currently `NULL`.

  It uses the same encryption rules as `up_for_v1_trigger/5`, so you should pass the
  exact same options used by the trigger.
  """
  @spec backfill_for_v1_column(String.t(), String.t(), String.t(), String.t(), keyword()) ::
          String.t()
  def backfill_for_v1_column(prefix, table, from, to, opts \\ []) when is_list(opts) do
    resolved_opts = resolve_cipher_options(prefix, table, from, to, opts)

    """
    UPDATE #{prefix}.#{table}
    SET #{to} =
      CASE
        WHEN #{from} IS NULL THEN NULL
        ELSE #{encrypted_value_sql(from, resolved_opts)}
      END
    WHERE #{backfill_condition_sql(to)};
    """
  end

  @doc """
  Returns SQL to drop a legacy trigger.

  ## ⚠️ Warning

  Dropping a trigger breaks encryption for the affected column pair.
  Only use this when intentionally removing or recreating the trigger (e.g., column rename).

  If recreating the trigger, you MUST use the exact same encryption parameters:
  - **time_bits**, **time_bucket**, **encrypt_time**: Same time configuration
  - **data_bits**: Same data bit size
  - **key**: Same encryption key
  - **rounds**: Same number of rounds
  - **functions_prefix**: Same schema where cipher functions reside

  For auto-generated keys, use `generate_key/4` with the **original** column names.

  ## Example

      # Column rename (seq -> sequence, id -> external_id)
      def change do
        execute FeistelCipher.down_for_legacy_trigger("public", "posts", "seq", "id")

        rename table(:posts), :seq, to: :sequence
        rename table(:posts), :id, to: :external_id

        old_key = FeistelCipher.generate_key("public", "posts", "seq", "id")

        execute FeistelCipher.up_for_legacy_trigger("public", "posts", "sequence", "external_id",
          time_bits: 15,
          time_bucket: 86400,
          data_bits: 38,
          key: old_key,
          rounds: 16,
          functions_prefix: "public"
        )
      end

  """
  @spec down_for_legacy_trigger(String.t(), String.t(), String.t(), String.t()) :: String.t()
  def down_for_legacy_trigger(prefix, table, from, to) do
    "DROP TRIGGER #{legacy_trigger_name(table, from, to)} ON #{prefix}.#{table};"
  end

  @doc """
  Same as `down_for_legacy_trigger/4` but targets v1 trigger naming convention.
  """
  @spec down_for_v1_trigger(String.t(), String.t(), String.t(), String.t()) :: String.t()
  def down_for_v1_trigger(prefix, table, from, to) do
    "DROP TRIGGER #{trigger_name(table, from, to)} ON #{prefix}.#{table};"
  end

  @doc """
  Generates a deterministic encryption key based on table/column information.

  Uses SHA-512 hash to derive a 31-bit key (valid range: 0 to 2^31-1).
  Same parameters always generate the same key, ensuring consistency across deployments.

  This is useful when recreating triggers (e.g., column rename) to maintain the same encryption key.

  ## Examples

      # Get the key used by the original trigger
      key = FeistelCipher.generate_key("public", "posts", "seq", "id")

      # Use it when recreating with new column names
      FeistelCipher.up_for_legacy_trigger("public", "posts", "sequence", "external_id", key: key)

  """
  @spec generate_key(String.t(), String.t(), String.t(), String.t()) :: non_neg_integer()
  def generate_key(prefix, table, from, to) do
    <<key::31, _::481>> = :crypto.hash(:sha512, "#{prefix}_#{table}_#{from}_#{to}")
    key
  end

  defp trigger_name(table, from, to) do
    "#{table}_encrypt_#{from}_to_#{to}_v1_trigger"
  end

  defp legacy_trigger_name(table, from, to) do
    "#{table}_encrypt_#{from}_to_#{to}_trigger"
  end

  defp encrypted_value_sql(from, resolved_opts) do
    data_component_sql =
      "#{resolved_opts.functions_prefix}.feistel_cipher_v1(#{from}, #{resolved_opts.data_bits}, #{resolved_opts.key}, #{resolved_opts.rounds})"

    "(#{time_component_sql(resolved_opts)} << #{resolved_opts.data_bits}) | #{data_component_sql}"
  end

  defp time_component_sql(%{time_bits: 0}) do
    "0"
  end

  defp time_component_sql(resolved_opts) do
    time_value_sql =
      "(floor((extract(epoch from now()) + #{resolved_opts.time_offset}::double precision) / #{resolved_opts.time_bucket}::double precision)::bigint & ((1::bigint << #{resolved_opts.time_bits}) - 1))"

    if resolved_opts.encrypt_time do
      "#{resolved_opts.functions_prefix}.feistel_cipher_v1(#{time_value_sql}, #{resolved_opts.time_bits}, #{resolved_opts.key}, #{resolved_opts.rounds})"
    else
      time_value_sql
    end
  end

  defp backfill_condition_sql(to) do
    "#{to} IS NULL OR #{to} = #{@backfill_sentinel}"
  end

  defp resolve_cipher_options(prefix, table, from, to, opts) do
    time_bits = Keyword.get(opts, :time_bits, 15)
    encrypt_time = Keyword.get(opts, :encrypt_time, false)
    use_time_prefix = time_bits > 0

    unless time_bits >= 0 do
      raise ArgumentError, "time_bits must be non-negative, got: #{time_bits}"
    end

    if encrypt_time and rem(time_bits, 2) != 0 do
      raise ArgumentError,
            "time_bits must be an even number when encrypt_time is true, got: #{time_bits}"
    end

    time_bucket = Keyword.get(opts, :time_bucket, 86400)
    time_offset = Keyword.get(opts, :time_offset, 0)

    if use_time_prefix do
      unless time_bucket > 0 do
        raise ArgumentError, "time_bucket must be positive, got: #{time_bucket}"
      end

      unless is_integer(time_offset) do
        raise ArgumentError, "time_offset must be an integer, got: #{inspect(time_offset)}"
      end
    end

    data_bits = Keyword.get(opts, :data_bits, 38)

    unless rem(data_bits, 2) == 0 do
      raise ArgumentError, "data_bits must be an even number, got: #{data_bits}"
    end

    unless data_bits >= 0 do
      raise ArgumentError, "data_bits must be non-negative, got: #{data_bits}"
    end

    max_total_bits = if encrypt_time, do: 62, else: 63

    unless time_bits + data_bits <= max_total_bits do
      raise ArgumentError,
            "time_bits + data_bits must be <= #{max_total_bits} when encrypt_time is #{encrypt_time}, got: #{time_bits} + #{data_bits} = #{time_bits + data_bits}"
    end

    rounds = Keyword.get(opts, :rounds, 16)

    unless rounds >= 1 and rounds <= 32 do
      raise ArgumentError, "rounds must be between 1 and 32, got: #{rounds}"
    end

    key = Keyword.get(opts, :key) || generate_key(prefix, table, from, to)
    validate_key!(key, "key")

    %{
      time_bits: time_bits,
      time_bucket: time_bucket,
      time_offset: time_offset,
      encrypt_time: encrypt_time,
      data_bits: data_bits,
      key: key,
      rounds: rounds,
      functions_prefix: Keyword.get(opts, :functions_prefix, "public")
    }
  end

  defp validate_key!(key, name) do
    max_key = Bitwise.bsl(1, 31) - 1

    unless key >= 0 and key <= max_key do
      raise ArgumentError,
            "#{name} must be between 0 and 2^31-1 (0..#{max_key}), got: #{key}"
    end
  end
end