# Shamir's Secret Sharing Scheme (SSSS) with m-of-n rule for 128-bit numbers.
# Author: Oleg Andreev <oleganza@gmail.com>
#
# * Deterministic, extensible algorithm: every combination of secret and threshold produces exactly the same shares on each run. More shares can be generated without invalidating the first ones.
# * This algorithm splits and restores 96/104/128-bit secrets with up to 16 shares and up to 16 shares threshold.
# * Secret is a binary 16-byte string below ffffffffffffffffffffffffffffff61.
# * Shares are 17-byte binary strings with first byte indicating threshold and share index (these are necessary for recovery).
# 
# See also: https://en.wikipedia.org/wiki/Shamir%27s_Secret_Sharing
require 'digest/sha2'
require 'securerandom'
module BTC
  class SecretSharing
    
    Order96  = 0xffffffffffffffffffffffef         # Largest prime below 2**96: (2**96 - 17)
    Order104 = 0xffffffffffffffffffffffffef       # Largest prime below 2**104: (2**104 - 17)
    Order128 = 0xffffffffffffffffffffffffffffff61 # Largest prime below 2**128: (2**128 - 159)
    
    def initialize(bitlength = 128)
      if bitlength == 128
        @order = Order128
      elsif bitlength == 104
        @order = Order104
      elsif bitlength == 96
        @order = Order96
      else
        raise ArgumentError, "Unsupported bit length #{bitlength}"
      end
      @bitlength = bitlength
    end
    
    def random
      be_from_int(SecureRandom.random_number(@order))
    end
  
    # Returns N strings, any M of them are enough to retrieve a secret.
    # Each string encodes X and Y coordinates and also M. X & M takes one byte, Y takes 16 bytes:
    # MMMMXXXX YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY YYYYYYYY
    def split(secret, m, n)
      prime = @order
      secret_num = int_from_be(secret)
      if secret_num >= @order
        raise "Secret cannot be encoded with #{@bitlength}-bit SSSS"
      end
      if !(n >= 1 && n <= 16)
        raise "N must be between 1 and 16"
      end
      if !(m >= 1 && m <= n)
        raise "M must be between 1 and N"
      end
      shares = []
      coef = [secret_num] # polynomial coefficients
      (m - 1).times do |i|
        # Generate unpredictable yet deterministic coefficients for each secret and M.
        coef << int_from_be(prng(secret + m.chr + i.chr))
      end
      n.times do |i|
        x = i + 1
        exp = 1
        y = coef[0]
        while exp < m
          y = (y + (coef[exp] * ((x**exp) % prime) % prime)) % prime
          exp += 1
        end
        # Encode share
        shares << string_from_point(m, x, y)
      end
      shares
    end
  
    # Transforms M 17-byte binary strings into original secret 16-byte binary string.
    # Each share string must be well-formed.
    def restore(shares)
      prime = @order
      shares = shares.dup.uniq
      raise "No shares provided" if shares.size == 0
      points = shares.map{|s| point_from_string(s) } # [[m,x,y],...]
      ms = points.map{|p| p[0]}.uniq
      xs = points.map{|p| p[1]}.uniq
      raise "Shares do not use the same M value" if ms.size > 1
      m = ms.first
      raise "All shares must have unique X values" if xs.size != points.size
      raise "Number of shares should be M or more" if points.size < m
      points = points[0, m] # make sure we have exactly M points
      y = 0
      points.size.times do |formula| # 0..(m-1)
        # Multiply the numerator across the top and denominators across the bottom to do Lagrange's interpolation
        numerator = 1
        denominator = 1
        points.size.times do |count| # 0..(m-1)
          if formula != count # skip element with i == j
            startposition = points[formula][1]
            nextposition = points[count][1]
            numerator = (numerator * -nextposition) % prime
            denominator = (denominator * (startposition - nextposition)) % prime
          end
        end
        value = points[formula][2]
        y = (prime + y + (value * numerator * modinv(denominator, prime))) % prime
      end
      return be_from_int(y)
    end

    def prng(seed)
      x = @order
      s = nil
      pad = "".b
      while x >= @order
        s = Digest::SHA2.digest(Digest::SHA2.digest(seed + pad))[0,@bitlength/8]
        x = int_from_be(s)
        pad = pad + "\x00".b
      end
      s
    end

    # Returns mmmmxxxx yyyyyyyy yyyyyyyy ... (16 bytes of y)
    def string_from_point(m, x, y)
      m = to_nibble(m)
      x = to_nibble(x)
      byte = [(m << 4) + x].pack("C")
      byte + be_from_int(y)
    end
  
    # returns [m, x, y]
    def point_from_string(s)
      byte = s.bytes.first
      m = from_nibble(byte >> 4)
      x = from_nibble(byte & 0x0f)
      y = int_from_be(s[1..-1])
      [m, x, y]
    end
  
    # Encodes values in range 1..16 to one nibble where all values are encoded as-is, 
    # except for 16 which becomes 0. This is to make strings look friendly for common cases when M,N < 16
    def to_nibble(x)
      x == 16 ? 0 : x
    end

    def from_nibble(x)
      x == 0 ? 16 : x
    end

    # Gives the decomposition of the gcd of a and b.  Returns [x,y,z] such that x = gcd(a,b) and y*a + z*b = x
    def gcd_decomposition(a,b)
      if b == 0
        [a, 1, 0]
      else
        n = a/b
        c = a % b
        r = gcd_decomposition(b,c)
        [r[0], r[2], r[1]-r[2]*n]
      end
    end

    # Gives the multiplicative inverse of k mod prime. In other words (k * modInverse(k)) % prime = 1 for all prime > k >= 1  
    def modinv(k, prime)
      k = k % prime
      r = (k < 0) ? -gcd_decomposition(prime,-k)[2] : gcd_decomposition(prime,k)[2]
      return (prime + r) % prime
    end

    def int_from_be(data)
      r = data.unpack("C*").reverse.inject({pos:0, total:0}) do |ctx, c|
        c = c << ctx[:pos]
        ctx[:pos] += 8
        ctx[:total] += c
        ctx
      end
      r[:total]
    end

    def be_from_int(i, pad = @bitlength)
      a = []
      while i > 0
        a.unshift(i % 256)
        i /= 256
      end
      a.pack("C*").rjust(pad / 8, "\x00")
    end

  end
end

if $0 == __FILE__
  
  ssss = BTC::SecretSharing.new(128)
  require_relative 'data.rb'
  
  # Usage 
  secret = ssss.random
  puts "Secret: #{BTC.to_hex(secret)}"
  shares = ssss.split(secret, 2, 3)
  shares.each do |share|
    puts "Share:  #{BTC.to_hex(share)}"
  end
  restored_secret = ssss.restore([shares[1], shares[0]])
  puts "Recovered secret with shares 2 and 1: #{BTC.to_hex(restored_secret)}"
  restored_secret = ssss.restore([shares[0], shares[2]])
  puts "Recovered secret with shares 1 and 3: #{BTC.to_hex(restored_secret)}"
  restored_secret = ssss.restore([shares[1], shares[2]])
  puts "Recovered secret with shares 2 and 3: #{BTC.to_hex(restored_secret)}"
  
  # Output:
  # Secret: d881c6f74ccac24997bb27040640a8eb
  # Share:  2147018841997da8d92211ad7590f85754
  # Share:  22b581498be6308f68ac6833e71bb0051e
  # Share:  2324010ad632e375f836beba58a667b387
  # Recovered secret with shares 2 and 1: d881c6f74ccac24997bb27040640a8eb
  # Recovered secret with shares 1 and 3: d881c6f74ccac24997bb27040640a8eb
  # Recovered secret with shares 2 and 3: d881c6f74ccac24997bb27040640a8eb
  
  # Test Vectors
  
  test_vectors = [
    {
      "secret" => "31415926535897932384626433832795",
      "1-of-1" => ["1131415926535897932384626433832795"],
      "1-of-2" => ["1131415926535897932384626433832795", "1231415926535897932384626433832795"],
      "2-of-2" => ["215af384f05d9b45f0e4e348f95b371acd", "2284a5b0ba67ddf44ea6422f8e82eb0e05"],
      "1-of-3" => ["1131415926535897932384626433832795", "1231415926535897932384626433832795", "1331415926535897932384626433832795"],
      "2-of-3" => ["215af384f05d9b45f0e4e348f95b371acd", "2284a5b0ba67ddf44ea6422f8e82eb0e05", "23ae57dc847220a2ac67a11623aa9f013d"],
      "3-of-3" => ["316cb005ab037e85ed9c8befbe72fef75c", "321387c8a1b34863197fae486ca60c1b97", "3325c8a20a62b62f16cceb6c6eccaa93a7"],
      "4-of-6" => ["416c4b3a8dc218696f8b1aed23385496eb", "429b14a744ce462bdc71b910b5cf0890ba", "4384d4d7881b01db3881cd0f17457112c8", 
                   "44f0c303944b6b73e265c52a42e9601a3c", "45a61663a602a2f238c80fa43408a7a57b", "466c062ff9e3c8529a531abee5f119b1ac"],
      "10-of-16"=>["a1a8b4077b75b0b18aefa63399d0b8d749", "a2e015e817190296d9ebe29f1c8cdc21c7", "a3c65760010c358c9760cece5da815edb4", "a4129891c5efd375a8367c854ab08010d6",
                   "a53c138386a55b0b35447ca03e44ab4eeb", "a6182993f21038c5d3bf548dac9dee7e20", "a769f010c04a4996b471a82addd4ea05d4", "a88e27a316dda9822f81616b2d48cb5e23",
                   "a9b0298820dc8c26989b6f8a2e8b00c3c4", "aa98042e1bcdf63b7283503ac4ad364380", "ab27bed0235b651dd92e764fa8cea25ba8", "ac05890d2177c48f4ec6cabd1047d9dbdc",
                   "adba7838775b82e4022af68f19d9985368", "aeb96045352c20fd24c6de8563cb2446f2", "af4f51af0a774592f9eabb71aaf0348def", "a06f50a680d22280f31b853d941c7eb158"],
    },
    {
      "secret" => "deadbeefcafebabedeadbeefcafebabe",
      "1-of-1" => ["11deadbeefcafebabedeadbeefcafebabe"],
      "2-of-2" => ["217f21b8a8329e69ea75a518485c8da19d", "221f95b2609a3e19160c9c71a0ee1c887c"],
      "2-of-3" => ["217f21b8a8329e69ea75a518485c8da19d", "221f95b2609a3e19160c9c71a0ee1c887c", "23c009ac1901ddc841a393caf97fab6ebc"],
      "3-of-3" => ["31d6b7c83a2587dd06be735c2ba5c719c0", "32762d76edcca00dd227bccb825a8daa75", "33bd0ecb0ac0474d211a8a0cf3e9526c3e"],
    },
    {
      "secret" => "ffffffffffffffffffffffffffffff60",
      "1-of-1" => ["11ffffffffffffffffffffffffffffff60"],
      "2-of-2" => ["21375c71bcaf077f5946f9e901efb9cf70", "226eb8e3795e0efeb28df3d203df739ee1"],
      "2-of-3" => ["21375c71bcaf077f5946f9e901efb9cf70", "226eb8e3795e0efeb28df3d203df739ee1", "23a61555360d167e0bd4edbb05cf2d6e52"],
      "3-of-3" => ["3112dac40bb910928263e5cf3971c39c8b", "32dec3f6359b1f7671aa60dd821c4969d3", "3363bb967da62cabcdd3712ad9ff916915"],
    },
    {
      "secret" => "00000000000000000000000000000000",
      "1-of-1" => ["1100000000000000000000000000000000"],
      "2-of-2" => ["2125df3f1da76af07c37689382bc8201a6", "224bbe7e3b4ed5e0f86ed127057904034c"],
      "2-of-3" => ["2125df3f1da76af07c37689382bc8201a6", "224bbe7e3b4ed5e0f86ed127057904034c", "23719dbd58f640d174a639ba88358604f2"],
      "3-of-3" => ["31651161eeddabb39134be97908f0d7d9e", "32671d1a7e6d7ef24037990a5285a75164", "33062329aeaf79bc0d088f5845e3cd7b52"],
    },
    {
      "secret" => "31415926535897932384626433", # 104 bits
      "1-of-1" => ["1131415926535897932384626433"],
      "1-of-2" => ["1131415926535897932384626433", "1231415926535897932384626433"],
      "2-of-2" => ["21a8453099fb8ae36aab2c1b6000", "221f49080da3bd2f4232d3d45bde"],
      "1-of-3" => ["1131415926535897932384626433", "1231415926535897932384626433", "1331415926535897932384626433"],
      "2-of-3" => ["21a8453099fb8ae36aab2c1b6000", "221f49080da3bd2f4232d3d45bde", "23964cdf814bef7b19ba7b8d57ab"],
      "3-of-3" => ["312b1880cc54fa4f009a4828f7d1", "3229a1cb3a58caea9c2bf0a31cb3", "332cdd38705eca6a65d87dd0d2d9"],
      "4-of-6" => ["414dc7bfe7e209630e44617d9737", "422da138cfafb6fa65f2556740c0", "43bd57245c5ccd2a2b2018645fa5", "44e972e30c89b7beeec062b9f2df", "459e7dd55ed6e28541c5ecacf956", "46c9015bd1e4b949b5236e8271e1"],
      "10-of-16"=>["a1b7a1ae5c3b1de94a8f4b8f7e05", "a25906c43a8969f6b5bf4be06006", "a3e97beb27daf0e2f81e0346a327", "a4d1a2cbee6be940edf213ccbc5f", "a50810074c61f554e9790453f951", "a629b4cf2ee2e5692af73337a0ca", "a7902bbc8ad42f8f9d15128e8758", "a8e6e7bce215a6f69d903a756a02", "a9834eca7f01f1a00215cf0ab33c", "aaa7b73a49f876d11a0b258fec2f", "abde0dc1d5af02f828197e7e0ac7", "acbc9befd7052707b1821140074e", "adadfaa8b410696a0d9089baaf84", "aebcb2502c16b5afd4f377c98514", "afeb62a04f2c15fa2f540d669b95", "a04b86c51029a5de8538e14edf9e"],
    },
    {
      "secret" => "deadbeefcafebabedeadbeefca",
      "1-of-1" => ["11deadbeefcafebabedeadbeefca"],
      "2-of-2" => ["210833159d705e79f32a1c8cea96", "2231b86c4b15be3927758b5ae551"],
      "2-of-3" => ["210833159d705e79f32a1c8cea96", "2231b86c4b15be3927758b5ae551", "235b3dc2f8bb1df85bc0fa28e00c"],
      "3-of-3" => ["3112462cb0571d40f6e704adf5b7", "320841c292cc488c414627a293e1", "33c0a080972a809c9dfc169cca48"],
    },
    {
      "secret" => "ffffffffffffffffffffffffee",
      "1-of-1" => ["11ffffffffffffffffffffffffee"],
      "2-of-2" => ["21b47b7bca3c91cfa72ec99d9ded", "2268f6f79479239f4e5d933b3bec"],
      "2-of-3" => ["21b47b7bca3c91cfa72ec99d9ded", "2268f6f79479239f4e5d933b3bec", "231d72735eb5b56ef58c5cd8d9eb"],
      "3-of-3" => ["31354043caf86f780c8d306fa32b", "32fc8f21351fc80c54a2b8604700", "3355ec983e7609bcd84097d1eba0"],
    },
    {
      "secret" => "00000000000000000000000000",
      "1-of-1" => ["1100000000000000000000000000"],
      "2-of-2" => ["219aa26f55d8a706cb6801023e74", "223544deabb14e0d96d002047cf9"],
      "2-of-3" => ["219aa26f55d8a706cb6801023e74", "223544deabb14e0d96d002047cf9", "23cfe74e0189f51462380306bb6d"],
      "3-of-3" => ["315a50b9d324cbf8cf4546d9e085", "32ca50d93a5f8028070814b77faa", "3350005e35b01c8da7486998dd80"],
    },
  ]
  
  test_vectors.each do |test|
    hexsecret = test.delete("secret")
    secret = BTC.from_hex(hexsecret)
    ssss = BTC::SecretSharing.new(secret.bytesize*8)
    test.each do |rule, defined_shares|
      m, n = rule.split("-of-").map{|x|x.to_i}
      puts "Testing #{hexsecret} #{rule}:"
      shares = ssss.split(secret, m, n)
      hexshares = shares.map{|s| BTC.to_hex(s)}
      failed = false
      if hexshares != defined_shares
        failed = true
        puts "Failed test:"
        puts "            Expected:  #{defined_shares.inspect}"
        puts "            Generated: #{hexshares.inspect}"
      end
      subshares = shares[0...m] # TODO: iterate over various combinations
      restored_secret = ssss.restore(subshares)
      if restored_secret != secret
        failed = true
        puts "Failed #{hexsecret} #{rule} test: failed to restore secret using #{subshares.map{|s| BTC.to_hex(s)}.inspect}"
      end
      if !failed
        puts "Ok."
      end
    end
  end
  
  
  begin
    require_relative 'base58'
    payload_length = 14 # 13 + checksum
    #payload_length = 15 # 14 + checksum
    lengths = {}
    (0..255).each do |ver|
      prefixes = []
      s = nil
      (0..255).each do |i|
        s = BTC::Base58.base58_from_data(ver.chr + (i.chr*payload_length))
        #lengths[s.length] = 1
        prefixes << s[0,1]
        s = BTC::Base58.base58_from_data(ver.chr + Digest::SHA2.digest(ver.to_s + i.to_s)[0,payload_length])
        lengths[s.length] = 1
        prefixes << s[0,1]
      end
      prefixes = prefixes.uniq
      puts   "Ver: #{ver} => #{prefixes}"
      if prefixes.size == 1
        puts "                    " + s
      end
    end
    
    puts "Lengths: #{lengths.keys.sort.inspect}"
  rescue => e
  end
end