{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0d84075-3705-4560-bddc-30fa35928a31",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mikolajlh/pqlattice/blob/main/docs/examples/LWE_dual_attack.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9903b6b1-144e-4907-b95c-c09e8b032c5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pqlattice as pq\n",
    "import numpy as np\n",
    "import math\n",
    "pq.settings.set_backend(\"fast\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "13b3e154-5276-4653-8883-1735b2beb067",
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_lwe_basis(A, q):\n",
    "    m, n = A.shape\n",
    "    Im = q * pq.as_integer(np.identity(m))\n",
    "    G = np.vstack((A.T, Im))\n",
    "    H, _ = pq.linalg.hnf(G)\n",
    "    rows, cols = H.shape\n",
    "    \n",
    "    return H[:m]\n",
    "\n",
    "def construct_sis_basis(A, q):\n",
    "    B_primal = construct_lwe_basis(A, q)\n",
    "    B_inv = np.linalg.inv(B_primal.astype(float))\n",
    "    B_dual = np.round(q * B_inv.T).astype(int)\n",
    "    return B_dual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e3e5aedb-df88-4edb-a992-394137833869",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20\n",
    "sigma = 2\n",
    "q = 1000\n",
    "m = 100\n",
    "secret_dist = \"ternary\"\n",
    "\n",
    "possible_values = []\n",
    "if secret_dist == \"binary\":\n",
    "    possible_values = [0, 1]\n",
    "elif secret_dist == \"ternary\":\n",
    "    possible_values = [-1, 0, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4e62260b-42af-43e8-96b3-090f08424d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "lwe = pq.random.LWE(n, q, sigma, secret_dist, 80)\n",
    "secret = lwe.secret\n",
    "A, b = lwe.sample_matrix(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2c7d7e24-de31-4e2b-ade3-f6da16e420ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recovering secret's components\n",
      "s[ 0]: Guessed  1 (Score: 0.9549) -> OK\n",
      "s[ 1]: Guessed  0 (Score: 0.9921) -> OK\n",
      "s[ 2]: Guessed  1 (Score: 0.9967) -> OK\n",
      "s[ 3]: Guessed  1 (Score: 0.9999) -> OK\n",
      "s[ 4]: Guessed  1 (Score: 1.0000) -> OK\n",
      "s[ 5]: Guessed  0 (Score: 0.9867) -> OK\n",
      "s[ 6]: Guessed  1 (Score: 0.9637) -> FAIL\n",
      "s[ 7]: Guessed  0 (Score: 1.0000) -> OK\n",
      "s[ 8]: Guessed  1 (Score: 0.9929) -> OK\n",
      "s[ 9]: Guessed  0 (Score: 1.0000) -> OK\n",
      "s[10]: Guessed  1 (Score: 0.9967) -> OK\n",
      "s[11]: Guessed  1 (Score: 0.9620) -> OK\n",
      "s[12]: Guessed  1 (Score: 0.9990) -> OK\n",
      "s[13]: Guessed  0 (Score: 0.9993) -> OK\n",
      "s[14]: Guessed  1 (Score: 0.9998) -> OK\n",
      "s[15]: Guessed  1 (Score: 0.9950) -> OK\n",
      "s[16]: Guessed -1 (Score: 0.9995) -> OK\n",
      "s[17]: Guessed  1 (Score: 0.9999) -> OK\n",
      "s[18]: Guessed  1 (Score: 0.9972) -> OK\n",
      "s[19]: Guessed  1 (Score: 0.9731) -> OK\n",
      "\n",
      "Real secret:\n",
      "[1 0 1 1 1 0 -1 0 1 0 1 1 1 0 1 1 -1 1 1 1]\n",
      "Recovered secret:\n",
      "[1 0 1 1 1 0 1 0 1 0 1 1 1 0 1 1 -1 1 1 1]\n",
      "\n",
      "accuracy: 95.00\n"
     ]
    }
   ],
   "source": [
    "recovered_secret = []\n",
    "\n",
    "print(f\"Recovering secret's components\")\n",
    "for i in range(n):\n",
    "    A_punctured = np.delete(A, i, axis=1)\n",
    "    G_dual = construct_sis_basis(A_punctured, q)\n",
    "    target_column = A[:, i]\n",
    "    v = pq.lattice.bkz(G_dual, 20)[0]\n",
    "    interaction = np.dot(v, target_column)\n",
    "    \n",
    "    best_guess = 0\n",
    "    max_score = -float('inf')\n",
    "    \n",
    "    projection = np.dot(v, b)\n",
    "    \n",
    "    for guess in possible_values:\n",
    "        \n",
    "        correction = guess * interaction\n",
    "        z = (projection - correction) % q\n",
    "        angle = (2 * math.pi * z) / q\n",
    "        score = math.cos(angle)\n",
    "        \n",
    "        if score > max_score:\n",
    "            max_score = score\n",
    "            best_guess = guess\n",
    "            \n",
    "    recovered_secret.append(best_guess)\n",
    "    \n",
    "    is_correct = (best_guess == secret[i])\n",
    "    status = \"OK\" if is_correct else \"FAIL\"\n",
    "    print(f\"s[{i:2}]: Guessed {best_guess:>2} (Score: {max_score:.4f}) -> {status}\")\n",
    "\n",
    "print()\n",
    "print(f\"Real secret:\")\n",
    "print(f\"{pq.as_integer(secret)}\")\n",
    "print(f\"Recovered secret:\")\n",
    "print(f\"{pq.as_integer(recovered_secret)}\")\n",
    "accuracy = (pq.as_integer(recovered_secret) == pq.as_integer(secret)).sum() / len(secret)\n",
    "print()\n",
    "print(f\"accuracy: {accuracy*100:.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
