ml_kem 0.1.0 → 0.1.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5d0131894b513113fb5492afe072bfec6ac2ba8d64531dbbd9c712cf419aff3a
-  data.tar.gz: f61a2aa77bef7d515d8b55c10dc66f07342b408d8884fbeca88e3d1fe673d4c1
+  metadata.gz: 291966b70102237cc2d2c36410dffafb2aef96f16381f0c4f25bb52b54e2666d
+  data.tar.gz: a12f02cda9894cd988451e38a6eebb4bb6cd0b889f1ae261a8fd64fc5cf3882d
 SHA512:
-  metadata.gz: 83a36606aa8c28521590dcd2c10a60f74b9e76e90562c60f2b1b3af188294567736c45ba128d40a322e63d65640461be8720004b7bd72a553f4db427655d9573
-  data.tar.gz: 1bf7d2be9236d916e4e5043f2377c630632a48277459e17ca521ba5890b61c54c31591c4fef6e86ab2bb22d12ee4c3e3cfefc9cd84c84953775cdf30746796bf
+  metadata.gz: 5b06d18dbf5a67425607e6cfadb21681d19ed0cee9efa8ad516e67874487fec7fd4ed93e079bbf923457824e309d4af8da9396019e8473c257210032c0b944d9
+  data.tar.gz: b683231323c1a52032acf6d084cf2e5ce2b1b508109b58e01ee340139091f16ba73915b7f2bbfc01a786922c7c65e32c7085ff9122b87426609467d577c1a259

data/lib/ml_kem/cli.rb

@@ -0,0 +1,147 @@
+# frozen_string_literal: true
+
+require 'thor'
+require 'ml_kem'
+require 'base64'
+
+module MLKEM
+  # Command Line Interface (CLI) for ML-KEM key encapsulation mechanism.
+  # Supports key generation, encapsulation, and decapsulation via Thor commands.
+  #
+  # @example Generate keys
+  #   $ mlkem keygen -p pub.pem -s priv.pem
+  #
+  # @example Encapsulate
+  #   $ mlkem encaps -p pub.pem -c ct.txt -k secret.key
+  #
+  # @example Decapsulate
+  #   $ mlkem decaps -s priv.pem -c ct.txt -k secret.key
+  #
+  # @since 0.1.0
+  class CLI < Thor
+    # Ensures the CLI exits with a non-zero status code on failure
+    def self.exit_on_failure?
+      true
+    end
+
+    # Global option for selecting the ML-KEM variant
+    class_option :variant,
+                 aliases: '-v',
+                 type: :string,
+                 default: 'ml_kem_768',
+                 desc: 'Select the ML-KEM variant (ml_kem_512, ml_kem_768, ml_kem_1024)'
+
+    desc "keygen", "Generate a key pair"
+    option :pk, aliases: "-p", type: :string, default: "public_key.pem", desc: "Output file for the public key"
+    option :sk, aliases: "-s", type: :string, default: "private_key.pem", desc: "Output file for the private key"
+    #
+    # Generates a public/private key pair and stores them in PEM format.
+    #
+    # @option options [String] :pk Output path for the public key.
+    # @option options [String] :sk Output path for the private key.
+    #
+    # @example
+    #   $ mlkem keygen -p pub.pem -s priv.pem
+    def keygen
+      mlkem = create_mlkem_instance
+      public_key, private_key = mlkem.keygen
+
+      File.write(options[:pk], encode_pem(public_key, "PUBLIC KEY"))
+      File.write(options[:sk], encode_pem(private_key, "PRIVATE KEY"))
+
+      puts "Keys generated:\n  Public Key: #{options[:pk]}\n  Private Key: #{options[:sk]}"
+    end
+
+    desc "encaps", "Encapsulate a secret using the public key"
+    option :pk, aliases: "-p", type: :string, required: true, desc: "Input file containing the public key (for encapsulation)"
+    option :sharedkey, aliases: "-k", type: :string, default: "shared_secret.key", desc: "Output file for the shared secret"
+    option :ciphertext, aliases: "-c", type: :string, default: "ciphertext.txt", desc: "Output file for the ciphertext"
+    #
+    # Encapsulates a shared secret using the given public key.
+    #
+    # @option options [String] :pk Path to PEM-encoded public key file.
+    # @option options [String] :sharedkey Output path for the base64-encoded shared secret.
+    # @option options [String] :ciphertext Output path for the base64-encoded ciphertext.
+    #
+    # @example
+    #   $ mlkem encaps -p pub.pem -c ct.txt -k secret.key
+    def encaps
+      mlkem = create_mlkem_instance
+      public_key_pem = File.read(options[:pk])
+      public_key = decode_pem(public_key_pem)
+
+      secret, ciphertext = mlkem.encaps(public_key)
+
+      File.write(options[:ciphertext], Base64.strict_encode64(ciphertext))
+      File.write(options[:sharedkey], Base64.strict_encode64(secret))
+
+      puts "Encapsulation complete:\n  Ciphertext: #{options[:ciphertext]}\n  Shared Secret: #{options[:sharedkey]}"
+    end
+
+    desc "decaps", "Decapsulate a ciphertext using the private key"
+    option :sk, aliases: "-s", type: :string, required: true, desc: "Input file containing the private key (for decapsulation)"
+    option :ciphertext, aliases: "-c", type: :string, required: true, desc: "Ciphertext file to decapsulate"
+    option :sharedkey, aliases: "-k", type: :string, default: "shared_secret.key", desc: "Output file for the shared secret"
+    #
+    # Decapsulates a ciphertext to recover the shared secret using the private key.
+    #
+    # @option options [String] :sk Path to PEM-encoded private key file.
+    # @option options [String] :ciphertext Path to base64-encoded ciphertext file.
+    # @option options [String] :sharedkey Output path for the base64-encoded shared secret.
+    #
+    # @example
+    #   $ mlkem decaps -s priv.pem -c ct.txt -k secret.key
+    def decaps
+      mlkem = create_mlkem_instance
+      private_key_pem = File.read(options[:sk])
+      private_key = decode_pem(private_key_pem)
+
+      ciphertext_base64 = File.read(options[:ciphertext])
+      ciphertext = Base64.decode64(ciphertext_base64)
+
+      secret = mlkem.decaps(private_key, ciphertext)
+
+      File.write(options[:sharedkey], Base64.strict_encode64(secret))
+
+      puts "Decapsulation complete. Shared secret written to #{options[:sharedkey]}"
+    end
+
+    private
+
+    # Initializes the ML-KEM instance using the provided variant.
+    #
+    # @return [MLKEM::MLKEM] Instance for encryption operations.
+    # @raise [ArgumentError] if the variant is not recognized.
+    def create_mlkem_instance
+      variant = options[:variant].to_sym
+      valid_variants = [:ml_kem_512, :ml_kem_768, :ml_kem_1024]
+
+      unless valid_variants.include?(variant)
+        raise ArgumentError, "Invalid variant: #{options[:variant]}. Valid options are: #{valid_variants.join(', ')}"
+      end
+
+      ::MLKEM::MLKEM.new(variant: variant)
+    end
+
+    # Decodes a PEM string into binary data.
+    #
+    # @param [String] pem_str The PEM-encoded key string.
+    # @return [String] Raw binary key data.
+    def decode_pem(pem_str)
+      base64_body = pem_str.lines.reject { |line|
+        line =~ /-----BEGIN|-----END/ || line.strip.empty?
+      }.join
+      Base64.decode64(base64_body)
+    end
+
+    # Encodes binary data into PEM format.
+    #
+    # @param [String] bin_str The binary data.
+    # @param [String] label The PEM label (e.g., "PUBLIC KEY").
+    # @return [String] PEM-formatted string.
+    def encode_pem(bin_str, label)
+      encoded = Base64.strict_encode64(bin_str).scan(/.{1,64}/).join("\n")
+      "-----BEGIN #{label}-----\n#{encoded}\n-----END #{label}-----\n"
+    end
+  end
+end

data/lib/ml_kem/constants.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+
+module MLKEM
+    # MLKEM::Constants module holds all static cryptographic parameters and tables
+    # used throughout the ML-KEM implementation.
+    #
+    # These include:
+    # - Approved parameter sets for ML-KEM (Table 2 of the specification)
+    # - Modulus (Q)
+    # - Polynomial degree (N)
+    # - Precomputed values for the Number-Theoretic Transform (NTT)
+    #
+    # @see https://csrc.nist.gov/pubs/fips/203/final NIST FIPS 203 - ML-KEM Specification
+    #
+    # @since 0.1.0
+    module Constants
+      ##
+      # Approved parameter sets for ML-KEM.
+      # These are tuples of the form [k, eta1, eta2, du, dv] where:
+      # - k    : number of polynomials in the public key
+      # - eta1 : noise sampling parameter for secret key
+      # - eta2 : noise sampling parameter for encryption
+      # - du   : number of bits revealed in the ciphertext part u
+      # - dv   : number of bits revealed in the ciphertext part v
+      #
+      # @return [Hash{String => Array<Integer>}]
+      PARAM_SETS = {
+        'ML-KEM-512'  => [2, 3, 2, 10, 4],
+        'ML-KEM-768'  => [3, 2, 2, 10, 4],
+        'ML-KEM-1024' => [4, 2, 2, 11, 5]
+      }.freeze
+  
+      ##
+      # Modulus used for all arithmetic operations in the ring Z_q.
+      #
+      # @return [Integer]
+      Q = 3329
+  
+      ##
+      # Degree of the polynomials used in ML-KEM, i.e., number of coefficients.
+      #
+      # @return [Integer]
+      N = 256
+  
+      ##
+      # Precomputed roots of unity for the Number-Theoretic Transform (NTT),
+      # as required by Appendix A of the ML-KEM specification.
+      #
+      # These are used for forward NTT operations.
+      #
+      # @return [Array<Integer>]
+      ZETA_NTT = [
+        1,    1729, 2580, 3289, 2642, 630,  1897, 848,
+        1062, 1919, 193,  797,  2786, 3260, 569,  1746,
+        296,  2447, 1339, 1476, 3046, 56,   2240, 1333,
+        1426, 2094, 535,  2882, 2393, 2879, 1974, 821,
+        289,  331,  3253, 1756, 1197, 2304, 2277, 2055,
+        650,  1977, 2513, 632,  2865, 33,   1320, 1915,
+        2319, 1435, 807,  452,  1438, 2868, 1534, 2402,
+        2647, 2617, 1481, 648,  2474, 3110, 1227, 910,
+        17,   2761, 583,  2649, 1637, 723,  2288, 1100,
+        1409, 2662, 3281, 233,  756,  2156, 3015, 3050,
+        1703, 1651, 2789, 1789, 1847, 952,  1461, 2687,
+        939,  2308, 2437, 2388, 733,  2337, 268,  641,
+        1584, 2298, 2037, 3220, 375,  2549, 2090, 1645,
+        1063, 319,  2773, 757,  2099, 561,  2466, 2594,
+        2804, 1092, 403,  1026, 1143, 2150, 2775, 886,
+        1722, 1212, 1874, 1029, 2110, 2935, 885,  2154
+      ].freeze
+  
+      ##
+      # Precomputed values used in point-wise multiplication during the NTT.
+      # These include both positive and negative roots of unity.
+      #
+      # @return [Array<Integer>]
+      ZETA_MUL = [
+        17, -17, 2761, -2761, 583, -583, 2649, -2649,
+        1637, -1637, 723, -723, 2288, -2288, 1100, -1100,
+        1409, -1409, 2662, -2662, 3281, -3281, 233, -233,
+        756, -756, 2156, -2156, 3015, -3015, 3050, -3050,
+        1703, -1703, 1651, -1651, 2789, -2789, 1789, -1789,
+        1847, -1847, 952, -952, 1461, -1461, 2687, -2687,
+        939, -939, 2308, -2308, 2437, -2437, 2388, -2388,
+        733, -733, 2337, -2337, 268, -268, 641, -641,
+        1584, -1584, 2298, -2298, 2037, -2037, 3220, -3220,
+        375, -375, 2549, -2549, 2090, -2090, 1645, -1645,
+        1063, -1063, 319, -319, 2773, -2773, 757, -757,
+        2099, -2099, 561, -561, 2466, -2466, 2594, -2594,
+        2804, -2804, 1092, -1092, 403, -403, 1026, -1026,
+        1143, -1143, 2150, -2150, 2775, -2775, 886, -886,
+        1722, -1722, 1212, -1212, 1874, -1874, 1029, -1029,
+        2110, -2110, 2935, -2935, 885, -885, 2154, -2154
+      ].freeze
+    end
+  end

data/lib/ml_kem/core/k_pke.rb

@@ -0,0 +1,171 @@
+# frozen_string_literal: true
+
+module MLKEM
+  # Core module containing the Public Key Encryption (K-PKE) implementation and the ML-KEM internal logic in two separate classes.
+  # 
+  # @since 0.1.0
+  module Core
+    # Implements the Public Key Encryption (K-PKE) as specified in ML-KEM (Kyber).
+    # Provides methods for key generation, encryption, and decryption using lattice-based NTT arithmetic.
+    #
+    # @since 0.1.0
+    class KPKE
+      # Initializes the KPKE engine with given parameters.
+      #
+      # @param [Integer] k       Security level (2, 3, or 4 for ML-KEM-512/768/1024)
+      # @param [Integer] eta1    Noise parameter for secret and error
+      # @param [Integer] eta2    Noise parameter for encryption errors
+      # @param [Integer] du      Compression bits for u
+      # @param [Integer] dv      Compression bits for v
+      # @param [Integer] q       Prime modulus
+      def initialize(k, eta1, eta2, du, dv, q = Constants::Q)
+        @k = k
+        @eta1 = eta1
+        @eta2 = eta2
+        @du = du
+        @dv = dv
+        @q = q
+
+        @poly_ops = Math::Polynomial.new(@q)
+        @ntt_ops = Math::NTT.new(@q)
+        @sampling = Math::Sampling.new(@q)
+      end
+
+      # Key generation algorithm (Algorithm 13).
+      #
+      # @param [String] d Random seed (32 bytes)
+      # @return [Array<String>] [public_key, secret_key] both as binary strings
+      def keygen(d)
+        rho, sig = Crypto::SymmetricPrimitives.g(d + [@k].pack('C'))
+        n = 0
+
+        a = Array.new(@k) { Array.new(@k) }
+        @k.times do |i|
+          @k.times do |j|
+            a[i][j] = @sampling.sample_ntt(rho + [j, i].pack('CC'))
+          end
+        end
+
+        s = Array.new(@k)
+        e = Array.new(@k)
+        @k.times do |i|
+          s[i] = @sampling.sample_poly_cbd(@eta1, Crypto::SymmetricPrimitives.prf(@eta1, sig, n))
+          n += 1
+          e[i] = @sampling.sample_poly_cbd(@eta1, Crypto::SymmetricPrimitives.prf(@eta1, sig, n))
+          n += 1
+        end
+
+        s.map! { |v| @ntt_ops.ntt(v) }
+        e.map! { |v| @ntt_ops.ntt(v) }
+
+        t = e.dup
+        @k.times do |i|
+          @k.times do |j|
+            t[i] = @poly_ops.add(t[i], @ntt_ops.multiply_ntts(a[i][j], s[j]))
+          end
+        end
+
+        ek_pke = Math::ByteOperations.byte_encode(12, t, @q) + rho
+        dk_pke = Math::ByteOperations.byte_encode(12, s, @q)
+
+        [ek_pke, dk_pke]
+      end
+
+      # Encryption algorithm (Algorithm 14).
+      #
+      # @param [String] ek_pke Public key
+      # @param [String] m      Message (32 bytes of encoded bits)
+      # @param [String] r      Randomness (32-byte seed)
+      # @return [String]       Ciphertext
+      def encrypt(ek_pke, m, r)
+        n = 0
+
+        t = Array.new(@k)
+        @k.times do |i|
+          t[i] = Math::ByteOperations.byte_decode(12, ek_pke[(384 * i)...(384 * (i + 1))], @q)
+        end
+        rho = ek_pke[(384 * @k)...(384 * @k + 32)]
+
+        a = Array.new(@k) { Array.new(@k) }
+        @k.times do |i|
+          @k.times do |j|
+            a[i][j] = @sampling.sample_ntt(rho + [j, i].pack('CC'))
+          end
+        end
+
+        y = Array.new(@k)
+        e1 = Array.new(@k)
+        @k.times do |i|
+          y[i] = @sampling.sample_poly_cbd(@eta1, Crypto::SymmetricPrimitives.prf(@eta1, r, n))
+          n += 1
+          e1[i] = @sampling.sample_poly_cbd(@eta2, Crypto::SymmetricPrimitives.prf(@eta2, r, n))
+          n += 1
+        end
+        e2 = @sampling.sample_poly_cbd(@eta2, Crypto::SymmetricPrimitives.prf(@eta2, r, n))
+
+        y.map! { |v| @ntt_ops.ntt(v) }
+
+        u = Array.new(@k) { Array.new(256, 0) }
+        @k.times do |i|
+          @k.times do |j|
+            u[i] = @poly_ops.add(u[i], @ntt_ops.multiply_ntts(a[j][i], y[j]))
+          end
+        end
+
+        @k.times do |i|
+          u[i] = @ntt_ops.inverse_ntt(u[i])
+          u[i] = @poly_ops.add(u[i], e1[i])
+        end
+
+        mu = @poly_ops.decompress(1, Math::ByteOperations.byte_decode(1, m, @q))
+
+        v = Array.new(256, 0)
+        @k.times do |i|
+          v = @poly_ops.add(v, @ntt_ops.multiply_ntts(t[i], y[i]))
+        end
+        v = @ntt_ops.inverse_ntt(v)
+        v = @poly_ops.add(v, e2)
+        v = @poly_ops.add(v, mu)
+
+        c1 = ''
+        @k.times do |i|
+          c1 += Math::ByteOperations.byte_encode(@du, @poly_ops.compress(@du, u[i]), @q)
+        end
+        c2 = Math::ByteOperations.byte_encode(@dv, @poly_ops.compress(@dv, v), @q)
+
+        c1 + c2
+      end
+
+      # Decryption algorithm (Algorithm 15).
+      #
+      # @param [String] dk_pke Secret key
+      # @param [String] c      Ciphertext
+      # @return [String]       Recovered message (32-byte encoded bitstring)
+      def decrypt(dk_pke, c)
+        c1 = c[0...(32 * @du * @k)]
+        c2 = c[(32 * @du * @k)...(32 * (@du * @k + @dv))]
+
+        up = Array.new(@k)
+        @k.times do |i|
+          up[i] = @poly_ops.decompress(@du,
+                  Math::ByteOperations.byte_decode(@du, c1[(32 * @du * i)...(32 * @du * (i + 1))], @q))
+        end
+
+        vp = @poly_ops.decompress(@dv, Math::ByteOperations.byte_decode(@dv, c2, @q))
+
+        s = Array.new(@k)
+        @k.times do |i|
+          s[i] = Math::ByteOperations.byte_decode(12, dk_pke[(384 * i)...(384 * (i + 1))], @q)
+        end
+
+        w = Array.new(256, 0)
+        @k.times do |i|
+          w = @poly_ops.add(w, @ntt_ops.multiply_ntts(s[i], @ntt_ops.ntt(up[i])))
+        end
+        w = @poly_ops.subtract(vp, @ntt_ops.inverse_ntt(w))
+
+        Math::ByteOperations.byte_encode(1, @poly_ops.compress(1, w), @q)
+      end
+    end
+  end
+end

data/lib/ml_kem/core/ml_kem_internal.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module MLKEM
+  module Core
+    # Internal implementation of ML-KEM (Kyber) encapsulation and decapsulation algorithms.
+    # Wraps K-PKE logic and adds hashing and key derivation per ML-KEM specification.
+    #
+    # Algorithms implemented:
+    # - Algorithm 16: ML-KEM.KeyGen_internal
+    # - Algorithm 17: ML-KEM.Encaps_internal
+    # - Algorithm 18: ML-KEM.Decaps_internal
+    #
+    # @since 0.1.0
+    class MLKEMInternal
+      # Constructs the internal ML-KEM engine.
+      #
+      # @param [Integer] k     Security level (1, 3, or 5)
+      # @param [Integer] eta1  Noise parameter η₁
+      # @param [Integer] eta2  Noise parameter η₂
+      # @param [Integer] du    Compression bits for u
+      # @param [Integer] dv    Compression bits for v
+      def initialize(k, eta1, eta2, du, dv)
+        @k = k
+        @eta1 = eta1
+        @eta2 = eta2
+        @du = du
+        @dv = dv
+        @q = Constants::Q
+        @kpke = KPKE.new(@k, @eta1, @eta2, @du, @dv, @q)
+      end
+
+      # ML-KEM Key Generation (Algorithm 16).
+      #
+      # @param [String] d A 32-byte seed for public key derivation.
+      # @param [String] z A 32-byte random secret value.
+      # @return [Array<String>] [ek, dk] Public and private key pair.
+      #
+      # @example
+      #   ek, dk = kem.keygen_internal(seed, z)
+      def keygen_internal(d, z)
+        ek_pke, dk_pke = @kpke.keygen(d)
+        ek = ek_pke
+        dk = dk_pke + ek + Crypto::SymmetricPrimitives.h(ek) + z
+        [ek, dk]
+      end
+
+      # ML-KEM Encapsulation (Algorithm 17).
+      #
+      # @param [String] ek Public key.
+      # @param [String] m A 32-byte message (uniformly random).
+      # @return [Array<String>] [k, c] Shared secret and ciphertext.
+      #
+      # @example
+      #   k, c = kem.encaps_internal(ek, m)
+      def encaps_internal(ek, m)
+        k, r = Crypto::SymmetricPrimitives.g(m + Crypto::SymmetricPrimitives.h(ek))
+        c = @kpke.encrypt(ek, m, r)
+        [k, c]
+      end
+
+      # ML-KEM Decapsulation (Algorithm 18).
+      #
+      # @param [String] dk Private key.
+      # @param [String] c  Ciphertext.
+      # @return [String]   Shared secret (32 bytes).
+      #
+      # @example
+      #   k = kem.decaps_internal(dk, c)
+      def decaps_internal(dk, c)
+        dk_pke = dk[0...(384 * @k)]
+        ek_pke = dk[(384 * @k)...(768 * @k + 32)]
+        h_val  = dk[(768 * @k + 32)...(768 * @k + 64)]
+        z      = dk[(768 * @k + 64)...(768 * @k + 96)]
+
+        mp = @kpke.decrypt(dk_pke, c)
+        kp, rp = Crypto::SymmetricPrimitives.g(mp + h_val)
+        kk = Crypto::SymmetricPrimitives.j(z + c)
+        cp = @kpke.encrypt(ek_pke, mp, rp)
+
+        c == cp ? kp : kk
+      end
+    end
+  end
+end

data/lib/ml_kem/crypto/hash_functions.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require 'sha3'
+
+module MLKEM
+  # Module containing an adjusted version of the Kyber hash functions
+  # used for cryptographic operations in ML-KEM.
+  #
+  # @since 0.1.0 
+  module Crypto
+    # Provides cryptographic hash functions used in ML-KEM (Kyber),
+    # including SHAKE128, SHAKE256, SHA3-256, and SHA3-512.
+    #
+    #
+    # @since 0.1.0
+    class HashFunctions
+      class << self
+        # Computes the SHAKE128 extendable-output function (XOF).
+        #
+        # @param [String] data Input byte string.
+        # @param [Integer] length Number of output bytes to produce.
+        # @return [String] XOF output of the specified length.
+        #
+        # @example
+        #   out = HashFunctions.shake128("seed", 32)
+        def shake128(data, length)
+          shake = SHA3::Digest::SHAKE_128.new
+          shake << data
+          shake.squeeze(length)
+        end
+
+        # Computes the SHAKE256 extendable-output function (XOF).
+        #
+        # @param [String] data Input byte string.
+        # @param [Integer] length Number of output bytes to produce.
+        # @return [String] XOF output of the specified length.
+        #
+        # @example
+        #   out = HashFunctions.shake256("key", 64)
+        def shake256(data, length)
+          shake = SHA3::Digest::SHAKE_256.new
+          shake << data
+          shake.squeeze(length)
+        end
+
+        # Computes the SHA3-256 hash of the given data.
+        #
+        # @param [String] data Input byte string.
+        # @return [String] 32-byte hash digest.
+        #
+        # @example
+        #   digest = HashFunctions.sha3_256("message")
+        def sha3_256(data)
+          digest = SHA3::Digest.new(:sha3_256)
+          digest.update(data)
+          digest.digest
+        end
+
+        # Computes the SHA3-512 hash of the given data.
+        #
+        # @param [String] data Input byte string.
+        # @return [String] 64-byte hash digest.
+        #
+        # @example
+        #   digest = HashFunctions.sha3_512("message")
+        def sha3_512(data)
+          digest = SHA3::Digest.new(:sha3_512)
+          digest.update(data)
+          digest.digest
+        end
+      end
+    end
+  end
+end

data/lib/ml_kem/crypto/symmetric_primitives.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require 'sha3'
+
+module MLKEM
+  module Crypto
+    # Provides symmetric cryptographic primitives as defined by ML-KEM (Kyber),
+    # including hash functions h, g, j and a pseudorandom function (PRF).
+    #
+    # These are used for key derivation, random generation, and message hashing.
+    #
+    # @since 0.1.0
+    class SymmetricPrimitives
+      class << self
+        # Hash function h: SHA3-256
+        #
+        # @param [String] x Input data to hash.
+        # @return [String] 32-byte digest.
+        #
+        # @example
+        #   digest = SymmetricPrimitives.h("message")
+        def h(x)
+          HashFunctions.sha3_256(x)
+        end
+
+        # Hash function g: SHA3-512, split into two 32-byte outputs.
+        #
+        # @param [String] x Input data to hash.
+        # @return [Array<String>] An array containing two 32-byte strings.
+        #
+        # @example
+        #   g1, g2 = SymmetricPrimitives.g("seed")
+        def g(x)
+          hash = HashFunctions.sha3_512(x)
+          [hash[0...32], hash[32...64]]
+        end
+
+        # Hash function j: SHAKE256 with 32-byte output.
+        #
+        # @param [String] s Input data to hash.
+        # @return [String] 32-byte XOF output.
+        #
+        # @example
+        #   result = SymmetricPrimitives.j("domain")
+        def j(s)
+          HashFunctions.shake256(s, 32)
+        end
+
+        # Pseudorandom Function (PRF)
+        #
+        # Uses SHAKE256 to expand `s || b` into `eta * 64` bytes of pseudorandom data.
+        #
+        # @param [Integer] eta Security parameter (e.g., 2 or 3).
+        # @param [String] s Secret seed.
+        # @param [Integer] b A single byte to include in domain separation.
+        # @return [String] Pseudorandom byte string of length `eta * 64`.
+        #
+        # @example
+        #   prf_output = SymmetricPrimitives.prf(3, "key", 0x01)
+        def prf(eta, s, b)
+          input = s + [b].pack('C*')
+          HashFunctions.shake256(input, eta * 64)
+        end
+      end
+    end
+  end
+end

data/lib/ml_kem/math/byte_operations.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module MLKEM
+  # Math module containing all the math related operations to NTT, byte operations, polynomial and sampling.
+  #
+  # @since 0.1.0
+  module Math
+    # Byte and bit manipulation operations used in the ML-KEM standard (2024).
+    #
+    # Provides methods to convert between bit arrays and byte strings,
+    # and to encode and decode arrays according to given parameters.
+    #
+    # @since 0.1.0
+    class ByteOperations
+      # Converts an array of bits into a byte string.
+      #
+      # Implements Algorithm 3, BitsToBytes(b).
+      #
+      # @param [Array<Integer>] b An array of bits (0 or 1).
+      # @return [String] The resulting byte string.
+      #
+      # @example
+      #   bits = [1,0,1,0,1,0,1,0]
+      #   ByteOperations.bits_to_bytes(bits) # => "\x55" (byte representation)
+      def self.bits_to_bytes(b)
+        l = b.length
+        a = Array.new(l / 8, 0)
+        (0...l).step(8) do |i|
+          x = 0
+          8.times do |j|
+            x += b[i + j] << j
+          end
+          a[i / 8] = x
+        end
+        a.pack('C*')
+      end
+
+      # Converts a byte string into an array of bits.
+      #
+      # Implements Algorithm 4, BytesToBits(B).
+      #
+      # @param [String] b A byte string.
+      # @return [Array<Integer>] The resulting array of bits (0 or 1).
+      #
+      # @example
+      #   bytes = "\x55"
+      #   ByteOperations.bytes_to_bits(bytes) # => [1,0,1,0,1,0,1,0]
+      def self.bytes_to_bits(b)
+        l = b.length
+        a = Array.new(8 * l, 0)
+        (0...(8 * l)).step(8) do |i|
+          x = b.bytes[i / 8]
+          8.times do |j|
+            a[i + j] = (x >> j) & 1
+          end
+        end
+        a
+      end
+
+      # Encodes an array of values into a byte string using dimension d and modulus q.
+      #
+      # Implements Algorithm 5, ByteEncode_d(F).
+      #
+      # @param [Integer] d Number of bits per value.
+      # @param [Array<Integer>, Array<Array<Integer>>] f Array of values or array of arrays to encode.
+      # @param [Integer] q Modulus for values (kept for compatibility; unused internally).
+      # @return [String] The encoded byte string.
+      #
+      # @example
+      #   d = 8
+      #   f = Array.new(256) { |i| i }
+      #   ByteOperations.byte_encode(d, f, 256) # => encoded string
+      #
+      # @example Encoding multiple arrays:
+      #   f_multi = [Array.new(256, 1), Array.new(256, 2)]
+      #   ByteOperations.byte_encode(d, f_multi, 256) # => concatenated encoded string
+      def self.byte_encode(d, f, q)
+        if f.first.is_a?(Array)
+          return f.map { |x| byte_encode(d, x, q) }.join
+        end
+
+        b = Array.new(256 * d, 0)
+
+        256.times do |i|
+          a = f[i]
+          d.times do |j|
+            b[d * i + j] = (a >> j) & 1
+          end
+        end
+        bits_to_bytes(b)
+      end
+
+      # Decodes a byte string into an array of values using dimension d and modulus q.
+      #
+      # Implements Algorithm 6, ByteDecode_d(B).
+      #
+      # @param [Integer] d Number of bits per value.
+      # @param [String] b Encoded byte string.
+      # @param [Integer] q Modulus for values.
+      # @return [Array<Integer>] The decoded array of values.
+      #
+      # @example
+      #   d = 8
+      #   encoded = ByteOperations.byte_encode(d, Array.new(256, 42), 256)
+      #   ByteOperations.byte_decode(d, encoded, 256) # => Array with value 42 repeated
+      def self.byte_decode(d, b, q)
+        bits_array = bytes_to_bits(b)
+        f = Array.new(256, 0)
+
+        256.times do |i|
+          a = 0
+          d.times do |j|
+            a += bits_array[d * i + j] << j
+          end
+          f[i] = a % q
+        end
+        f
+      end
+    end
+  end
+end

data/lib/ml_kem/math/ntt.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+module MLKEM
+  module Math
+    # Implements the Number Theoretic Transform (NTT) and related operations
+    # as defined in the ML-KEM standard (2024), used for polynomial arithmetic
+    # in the ring Z_q[X]/(X^256 + 1).
+    #
+    # Includes forward and inverse NTT transforms, pointwise multiplication,
+    # and the base case multiplication algorithm.
+    #
+    # @since 0.1.0
+    class NTT
+      # Initializes an NTT instance with a modulus `q`.
+      #
+      # @param [Integer] q The modulus for all modular arithmetic.
+      #
+      # @example
+      #   ntt = NTT.new
+      def initialize(q = Constants::Q)
+        @q = q
+      end
+
+      # Applies the forward Number Theoretic Transform to a polynomial.
+      #
+      # Implements Algorithm 9, NTT(f).
+      #
+      # @param [Array<Integer>] f A 256-element array representing the polynomial.
+      # @return [Array<Integer>] The transformed polynomial.
+      #
+      # @example
+      #   transformed = ntt.ntt(original_poly)
+      def ntt(f)
+        f = f.dup
+        i = 1
+        le = 128
+        
+        while le >= 2
+          (0...256).step(2 * le) do |st|
+            ze = Constants::ZETA_NTT[i]
+            i += 1
+            (st...(st + le)).each do |j|
+              t = (ze * f[j + le]) % @q
+              f[j + le] = (f[j] - t) % @q
+              f[j] = (f[j] + t) % @q
+            end
+          end
+          le /= 2
+        end
+        
+        f
+      end
+
+      # Applies the inverse Number Theoretic Transform to a polynomial.
+      #
+      # Implements Algorithm 10, NTT^<sup>-1</sup>(f).
+      #
+      # @param [Array<Integer>] f A 256-element array in the NTT domain.
+      # @return [Array<Integer>] The inverse-transformed polynomial.
+      #
+      # @example
+      #   original = ntt.inverse_ntt(transformed_poly)
+      def inverse_ntt(f)
+        f = f.dup
+        i = 127
+        le = 2
+        
+        while le <= 128
+          (0...256).step(2 * le) do |st|
+            ze = Constants::ZETA_NTT[i]
+            i -= 1
+            (st...(st + le)).each do |j|
+              t = f[j]
+              f[j] = (t + f[j + le]) % @q
+              f[j + le] = (ze * (f[j + le] - t)) % @q
+            end
+          end
+          le *= 2
+        end
+        
+        f.map { |x| (x * 3303) % @q }
+      end
+
+      # Performs pointwise multiplication in the NTT domain.
+      #
+      # Implements Algorithm 11, MultiplyNTTs(~f, ~g).
+      #
+      # @param [Array<Integer>] f First polynomial in NTT form.
+      # @param [Array<Integer>] g Second polynomial in NTT form.
+      # @return [Array<Integer>] Result of the pointwise multiplication.
+      #
+      # @example
+      #   product = ntt.multiply_ntts(ntt_f, ntt_g)
+      def multiply_ntts(f, g)
+        h = []
+        (0...256).step(2) do |ii|
+          h.concat(base_case_multiply(f[ii], f[ii + 1], g[ii], g[ii + 1], 
+                                      Constants::ZETA_MUL[ii / 2]))
+        end
+        h
+      end
+
+      # Performs the base case multiplication for two polynomial coefficient pairs.
+      #
+      # Implements Algorithm 12, BaseCaseMultiply(a0, a1, b0, b1, gamma).
+      #
+      # @param [Integer] a0 Coefficient a_0
+      # @param [Integer] a1 Coefficient a_1
+      # @param [Integer] b0 Coefficient b_0
+      # @param [Integer] b1 Coefficient b_1
+      # @param [Integer] gam Precomputed gamma value
+      # @return [Array<Integer>] The resulting coefficients [c0, c1]
+      #
+      # @example
+      #   ntt.base_case_multiply(3, 4, 5, 6, gamma) # => [expected_c0, expected_c1]
+      def base_case_multiply(a0, a1, b0, b1, gam)
+        c0 = (a0 * b0 + a1 * b1 * gam) % @q
+        c1 = (a0 * b1 + a1 * b0) % @q
+        [c0, c1]
+      end
+    end
+  end
+end

data/lib/ml_kem/math/polynomial.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module MLKEM
+  module Math
+    # Implements polynomial arithmetic and compression/decompression
+    # operations used in ML-KEM (Kyber) cryptographic schemes.
+    #
+    # Provides basic modular operations such as addition and subtraction
+    # as well as lossy compression methods used to reduce bandwidth.
+    #
+    # @since 0.1.0
+    class Polynomial
+      # Initializes a Polynomial instance with a modulus q.
+      #
+      # @param [Integer] q The modulus used for polynomial arithmetic.
+      #
+      # @example
+      #   poly = Polynomial.new
+      def initialize(q = Constants::Q)
+        @q = q
+      end
+
+      # Adds two polynomials coefficient-wise modulo q.
+      #
+      # @param [Array<Integer>] f First polynomial.
+      # @param [Array<Integer>] g Second polynomial.
+      # @return [Array<Integer>] Resulting polynomial (f + g) mod q.
+      #
+      # @example
+      #   result = poly.add([1, 2], [3, 4]) # => [4, 6]
+      def add(f, g)
+        f.zip(g).map { |a, b| (a + b) % @q }
+      end
+
+      # Subtracts one polynomial from another coefficient-wise modulo q.
+      #
+      # @param [Array<Integer>] f Minuend polynomial.
+      # @param [Array<Integer>] g Subtrahend polynomial.
+      # @return [Array<Integer>] Resulting polynomial (f - g) mod q.
+      #
+      # @example
+      #   result = poly.subtract([5, 3], [2, 1]) # => [3, 2]
+      def subtract(f, g)
+        f.zip(g).map { |a, b| (a - b) % @q }
+      end
+
+      # Compresses the coefficients of a polynomial to `d` bits.
+      #
+      # Lossy operation used to reduce size during transmission.
+      #
+      # @param [Integer] d Number of bits to compress to.
+      # @param [Array<Integer>] xv Polynomial coefficients.
+      # @return [Array<Integer>] Compressed coefficients.
+      #
+      # @example
+      #   compressed = poly.compress(4, [0, 1000, 2000])
+      def compress(d, xv)
+        xv.map do |x|
+          (((x << d) + (@q - 1) / 2) / @q) % (1 << d)
+        end
+      end
+
+      # Decompresses `d`-bit values back into approximate polynomial coefficients.
+      #
+      # Inverse of `#compress`, though lossy.
+      #
+      # @param [Integer] d Number of bits the coefficients were compressed to.
+      # @param [Array<Integer>] yv Compressed coefficients.
+      # @return [Array<Integer>] Approximate original coefficients.
+      #
+      # @example
+      #   decompressed = poly.decompress(4, [0, 5, 10])
+      def decompress(d, yv)
+        yv.map do |y|
+          (@q * y + (1 << (d - 1))) >> d
+        end
+      end
+    end
+  end
+end

data/lib/ml_kem/math/sampling.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require 'sha3'
+
+module MLKEM
+  module Math
+    # Implements sampling algorithms used in ML-KEM (Kyber),
+    # including NTT-domain sampling and centered binomial distribution (CBD).
+    #
+    # These routines are used to generate polynomial coefficients
+    # from uniformly random byte arrays or XOF outputs.
+    #
+    # @since 0.1.0
+    class Sampling
+      # Initializes the Sampling object with a modulus `q`.
+      #
+      # @param [Integer] q The modulus used in coefficient arithmetic.
+      #
+      # @example
+      #   sampling = Sampling.new
+      def initialize(q = Constants::Q)
+        @q = q
+      end
+
+      # Samples a polynomial in the NTT domain from a byte string using SHAKE128.
+      #
+      # Implements Algorithm 7, SampleNTT(B).
+      #
+      # This uses rejection sampling to select values uniformly < q from SHAKE128 output.
+      #
+      # @param [String] b Input byte string to expand into polynomial coefficients.
+      # @return [Array<Integer>] A 256-element array of sampled coefficients < q.
+      #
+      # @example
+      #   sampled = sampling.sample_ntt(seed_bytes)
+      def sample_ntt(b)
+        xof_data = Crypto::HashFunctions.shake128(b, 1024)
+        j = 0
+        a = []
+        i = 0
+        
+        while j < 256 && i < xof_data.length - 2
+          c = xof_data.bytes[i, 3]
+          d1 = c[0] + 256 * (c[1] % 16)
+          d2 = (c[1] / 16) + 16 * c[2]
+          
+          if d1 < @q
+            a << d1
+            j += 1
+          end
+          
+          if d2 < @q && j < 256
+            a << d2
+            j += 1
+          end
+          
+          i += 3
+        end
+        
+        a
+      end
+
+      # Samples a polynomial using centered binomial distribution (CBD).
+      #
+      # Implements Algorithm 8, SamplePolyCBD_eta(B).
+      #
+      # The result is a noise polynomial used in ML-KEM.
+      #
+      # @param [Integer] eta CBD parameter (usually 2 or 3).
+      # @param [String] b Byte string used as input entropy.
+      # @return [Array<Integer>] A 256-element array of polynomial coefficients modulo q.
+      #
+      # @example
+      #   noise = sampling.sample_poly_cbd(3, seed_bytes)
+      def sample_poly_cbd(eta, b)
+        b_bits = Math::ByteOperations.bytes_to_bits(b)
+        f = Array.new(256, 0)
+        
+        256.times do |i|
+          x = b_bits[(2 * i * eta)...((2 * i + 1) * eta)].sum
+          y = b_bits[((2 * i + 1) * eta)...((2 * i + 2) * eta)].sum
+          f[i] = (x - y) % @q
+        end
+        
+        f
+      end
+    end
+  end
+end

data/lib/ml_kem/version.rb

@@ -6,5 +6,5 @@
 # @author MarioRgzLpz <https://github.com/MarioRgzLpz>
 # @since 0.1.0
 module MLKEM
-  VERSION = "0.1.0"
+  VERSION = "0.1.1"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ml_kem
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.1
 platform: ruby
 authors:
 - MarioRgzLpz
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2025-07-03 00:00:00.000000000 Z
+date: 2025-07-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: sha3
@@ -30,56 +30,70 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.3'
+        version: 1.3.2
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.3'
+        version: 1.3.2
 - !ruby/object:Gem::Dependency
   name: yard
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.9'
+        version: 0.9.37
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.9'
+        version: 0.9.37
 - !ruby/object:Gem::Dependency
   name: minitest
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.22'
+        version: 5.25.5
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.22'
+        version: 5.25.5
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '13.3'
+        version: 13.3.0
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '13.3'
+        version: 13.3.0
+- !ruby/object:Gem::Dependency
+  name: irb
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 1.15.2
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 1.15.2
 description: A Ruby gem providing an implementation of the ML-KEM (formerly Kyber)
   key-encapsulation mechanism for post-quantum cryptography standards.
 email:
@@ -91,11 +105,19 @@ extra_rdoc_files: []
 files:
 - LICENSE.txt
 - README.md
-- Rakefile
 - exe/mlkem
 - lib/ml_kem.rb
+- lib/ml_kem/cli.rb
+- lib/ml_kem/constants.rb
+- lib/ml_kem/core/k_pke.rb
+- lib/ml_kem/core/ml_kem_internal.rb
+- lib/ml_kem/crypto/hash_functions.rb
+- lib/ml_kem/crypto/symmetric_primitives.rb
+- lib/ml_kem/math/byte_operations.rb
+- lib/ml_kem/math/ntt.rb
+- lib/ml_kem/math/polynomial.rb
+- lib/ml_kem/math/sampling.rb
 - lib/ml_kem/version.rb
-- sig/ml_kem.rbs
 homepage: https://github.com/MarioRgzLpz/ml_kem
 licenses:
 - MIT

data/Rakefile

@@ -1,8 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/gem_tasks"
-require "minitest/test_task"
-
-Minitest::TestTask.create
-
-task default: :test

data/sig/ml_kem.rbs

@@ -1,4 +0,0 @@
-module MlKem
-  VERSION: String
-  # See the writing guide of rbs: https://github.com/ruby/rbs#guides
-end