z3-shamir-secret-sharing.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "z3-shamir-secret-sharing.ipynb",
      "version": "0.3.2",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "metadata": {
        "id": "v-nPb_dQmluX",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "bdb28f64-f5a9-4a8d-8d94-42fe03f6afd5"
      },
      "cell_type": "code",
      "source": [
        "!pip install z3-solver"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: z3-solver in /usr/local/lib/python3.6/dist-packages (4.8.0.0.post1)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "CX6NfTeIuRVR",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Definition of simple functions to encode a string into an integer and back\n",
        "\n",
        "import binascii\n",
        "\n",
        "def str_to_int(str):\n",
        "  byts = str.encode(\"utf-8\")\n",
        "  return int(binascii.hexlify(byts), 16)\n",
        "\n",
        "def int_to_str(i):\n",
        "  hx = hex(i)[2:]\n",
        "  secret = binascii.unhexlify(hx)\n",
        "  return secret.decode(\"utf-8\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "cEFJK3hqzBIw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "# Function that generates another function given a list of coefficient.\n",
        "# (also determining the polynomial's degree)\n",
        "# https://en.wikipedia.org/wiki/Shamir%27s_Secret_Sharing#Usage\n",
        "def generate_polynomial_fn(a_vars):\n",
        "  parts = []\n",
        "  for degree in range(len(a_vars)):\n",
        "    # closure on variable 'degree' using default parameter\n",
        "    def polynomial_part(x, d=degree):\n",
        "      return a_vars[d] * x**d\n",
        "    parts.append(polynomial_part)\n",
        "    \n",
        "  # Returns a lambda that sums all part of the polynomial for a given 'x'\n",
        "  return lambda x: sum(map(lambda f: f(x), parts))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "FirYCvc_QBv-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def shamir_split(m, n, secret):\n",
        "  # Choosing random numbers that are not too far away from our secret\n",
        "  rnds = [int(random.uniform(0.2, 0.8) * secret) for _ in range(m - 1)]\n",
        "  f = generate_polynomial_fn([secret] + rnds)\n",
        "  # returns the evaluation of the polynomial function for 0 < x < n+1\n",
        "  # WARNING: here x=0 would reveal the secret (f(0) == secret)\n",
        "  return [(x, f(x)) for x in range(1, n + 1)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "LhHMr-2RI54t",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from z3 import *\n",
        "\n",
        "# https://en.wikipedia.org/wiki/Shamir%27s_Secret_Sharing#Solution\n",
        "def shamir_resolve(splits):\n",
        "  degree = len(splits)\n",
        "  \n",
        "  # Z3 way of declaring unknown variables\n",
        "  a = [Int(f'a{i}') for i in range(0, degree)]\n",
        "  solver = Solver()\n",
        "  \n",
        "  # Create a simple system of equations from the splits\n",
        "  for x, y in splits:\n",
        "    p = generate_polynomial_fn(a)\n",
        "    solver.add(y == p(x))\n",
        "  \n",
        "  # This part is Z3 specific, we check and solve the system of equations\n",
        "  solver.check()\n",
        "  model = solver.model()\n",
        "  sec_int = model[a[0]].as_long()\n",
        "  \n",
        "  return int_to_str(sec_int)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "5CYyGBqAtcYQ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "secret = \"this is so secret\"\n",
        "\n",
        "# We encode our secret into an integer\n",
        "sec_int = str_to_int(secret)\n",
        "\n",
        "# Make sure our reverse function works as expected\n",
        "assert(int_to_str(sec_int) == secret)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "rgsF-5tQ0dM9",
        "colab_type": "code",
        "outputId": "6c57e9e9-bdb0-438c-cdad-5b5991942729",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "cell_type": "code",
      "source": [
        "m = 7\n",
        "n = 12\n",
        "splits = shamir_split(m, n, sec_int)\n",
        "\n",
        "assert(len(splits) == n)\n",
        "print(splits)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[(1, 128578265812178243783178166342781756401012), (2, 1929330870138865244530456899917352686216564), (3, 15336725934380767844321927151285452194080116), (4, 71837137908574446326760744108841065260410228), (5, 244273538839691018525711883516500403855451508), (6, 673035469417389632886173932031131345992443252), (7, 1598673901339026316157179490230839621761262964), (8, 3400940990993918195716724188274111746342544756), (9, 6644254724466863096528724314209814700002272628), (10, 12129588453860914512732003054938052355066848628), (11, 20952785324939411953861305349821878649878635892), (12, 34569297596087266665700341356949867509731976564)]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "9qA34g-K0fnD",
        "colab_type": "code",
        "outputId": "9286d2d4-3c05-46fc-e519-64b21d58a99c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "cell_type": "code",
      "source": [
        "# Picking arbitrarily a subset of the splits\n",
        "split_subset = splits[3:10]\n",
        "\n",
        "# Make sure the chosen subset is greater or equal to m\n",
        "assert(len(split_subset) >= m)\n",
        "\n",
        "# Reconstruct the secret using only the splits (and the 'x' value)\n",
        "shamir_resolve(split_subset)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'this is so secret'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    }
  ]
}