{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = random.getrandbits(1024)\n",
    "#Just choosing a random 1024 bit number"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = random.getrandbits(1008)\n",
    "m = 2<<1008\n",
    "m = m + random.getrandbits(1008)\n",
    "#This is the hidden message we wish to find."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = 2**1008\n",
    "B2 = 2*B\n",
    "B3 = 3*B\n",
    "ocalls=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ceildiv(x,y):   #ceildiv(a,b) = ceil(a/b)\n",
    "    return x//y + (x%y != 0)\n",
    "def floordiv(x,y):  #floordiv(a,b) = floor(a/b)\n",
    "    return x//y\n",
    "\n",
    "def padCheck(y):\n",
    "    global ocalls\n",
    "    ocalls+=1\n",
    "    if y >= B2 and y < B3:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "#The usual PKSC#1 padding does more checks (there must be at least 8 padded bytes etc.)\n",
    "#This is a simpler version, that also increase the likelihood of hitting a properly padded string"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding an s\n",
    "\n",
    "The block `findS(smin)` starts from `smin`, and keeps trying out various `s`'s until it finds one that makes the `padCheck(m*s)` accept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findS(smin):\n",
    "    global m\n",
    "    global n\n",
    "    si = smin\n",
    "    while(True):\n",
    "        si+=1\n",
    "        mi = (m*si)%n\n",
    "        if padCheck(mi):\n",
    "            return si\n",
    "            break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figuring out the range of r's\n",
    "\n",
    "Once we have found an $s$, and if we are assuming that $a \\leq m \\leq b$, then we know that $2B \\leq ms - rn \\leq 3B-1$ for some positive integer $r$. This gives the bound\n",
    "$$\n",
    "\\frac{bs - 2B}{n} \\geq \\frac{ms - 2B}{n} \\geq r \\geq \\frac{ms-3B+1}{n} \\geq \\frac{as - 3B+1}{n}\n",
    "$$\n",
    "The function `findrRanges(s,a,b)` returns the lower and upper bound based on the above formula."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findrRanges(s,a,b):\n",
    "    rmin = ceildiv((a*s)-B3+1,n)\n",
    "    rmax = floordiv((b*s) - B2,n)\n",
    "    return(rmin,rmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## An optimisation in case of a single interval\n",
    "\n",
    "In [Bleichenbacher's original paper](http://archiv.infsec.ethz.ch/education/fs08/secsem/bleichenbacher98.pdf), he performs an optimisation in the setting when we have a single interval that is not the initial inverval of $(2B, 3B-1)$. In this case, you find $r \\geq \\frac{2(bs - 2B)}{n}$, and look for an $s$ in the range\n",
    "$$\n",
    "\\frac{2B + rn}{b} \\leq s \\leq \\frac{3B-1 +rn}{a}.\n",
    "$$\n",
    "Use this $s$ and proceed. \n",
    "\n",
    "The idea is that this results in a new interval that's no more than half the original interval. You can see the details in [Bleichenbacher's paper](http://archiv.infsec.ethz.ch/education/fs08/secsem/bleichenbacher98.pdf). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findS_opt(s,a,b):\n",
    "    global m\n",
    "    global n\n",
    "    global B2\n",
    "    global B3\n",
    "    \n",
    "    r = floordiv(2*((b*s) - B2),n)\n",
    "    while(True):\n",
    "        found = False\n",
    "        while(True):\n",
    "            for si in range(ceildiv(B2+(r*n),b),floordiv(B3-1+(r*n),a)+1):\n",
    "                mi = (si*m)%n\n",
    "                if padCheck(mi):\n",
    "                    found=True\n",
    "                    return(si)\n",
    "            if not found:\n",
    "                r+=1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The attack\n",
    "\n",
    "If we already knew that $a \\leq m \\leq b$ and that $2B \\leq ms-rn \\leq 3B$, then we get that\n",
    "$$\n",
    "\\frac{2B + rn}{s} \\leq m \\leq \\frac{3B + rn -1}{s}, \n",
    "$$\n",
    "and we can intersect this with the old interval $(a,b)$. The code below basically find an $s$ and updates the interval, until it ends up with a single interval of length $1$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Attack():\n",
    "    global ocalls,n,B2,B3,m\n",
    "    si = ceildiv(n,B3)-1\n",
    "    newM = set([])\n",
    "    newM.add((B2,B3-1))\n",
    "    ocalls=0\n",
    "    while(True):\n",
    "        # First find an si\n",
    "        if(len(newM)>1 or (B2,B3-1) in newM):\n",
    "            si = findS(si)\n",
    "        elif(len(newM)==1): #use the optimised version in the special case\n",
    "            (a,b) = newM.pop()\n",
    "            newM.add((a,b))\n",
    "            si = findS_opt(si,a,b)\n",
    "        \n",
    "        #update the intervals\n",
    "        newMM = set([])\n",
    "        for (a,b) in newM:\n",
    "            (r1,r2) = findrRanges(si,a,b)\n",
    "            for r in range(r1,r2+1):\n",
    "                aa = ceildiv(B2 + (r*n),si)\n",
    "                bb = floordiv(B3 - 1 + (r*n),si)\n",
    "                newa = max(a,aa)\n",
    "                newb = min(b,bb)\n",
    "                if newa <= newb:\n",
    "                    newMM.add((newa,newb))\n",
    "        if len(newMM)>0:\n",
    "            newM = newMM\n",
    "        else:\n",
    "            print(\"Something went wrong!\")\n",
    "            exit(-1)\n",
    "        if len(newM) == 1:\n",
    "            (a,b)= newM.pop()\n",
    "            newM.add((a,b))\n",
    "            if a==b:\n",
    "                print(\"Oracle calls:\", ocalls)\n",
    "                return a\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Oracle calls: 4506\n",
      "Found m:  5971331115703357078891961541342638949664044605888442834663177475828153859701294020291819323895945632313357163311109747139028146977957416417384602481528503116951057112305786705469189482505678846379821056382535831302517472639355811058140760853748105564413984688674138740815451264447663254207534472088145088\n",
      "CPU times: user 69.5 ms, sys: 3.83 ms, total: 73.3 ms\n",
      "Wall time: 206 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "guess = Attack()\n",
    "if m==guess:\n",
    "    print(\"Found m: \",m)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}