keyrecovery_plot.py

#!/usr/bin/env python3

WITH_SAGE = True
# usage:
# sage -python keyrecovery_plot.py -n 5 > warp.tex && latexmk warp.tex

import sys
import argparse
from math import log

if WITH_SAGE:
    from sage.all import Matrix, GF

pi = [31, 6, 29, 14, 1, 12, 21, 8, 27, 2, 3, 0, 25, 4, 23, 10, 15, 22, 13, 30, 17, 28, 5, 24, 11, 18, 19, 16, 9, 20, 7, 26]


### TEX WRAPPER ###

def tex_init(options="", standalone=False):
    print(r"% generated by " + options)
    if standalone:
        print(r"\documentclass[11pt]{standalone}")
    else:
        print(r"\documentclass[a4paper,11pt]{scrartcl}")
    print(r"""
\usepackage{warp}
\usepackage[margin=1in]{geometry}
\usetikzlibrary{graphs,graphdrawing,quotes}
\usegdlibrary{trees}

\begin{document}
""", end="")

def tex_fin(rounds):
    print(r"""
\end{document}
""")

def tex_figure_init(standalone=False):
    if not standalone:
        print(r"""
\begin{figure}[htp!]
""", end="")

def tex_figure_fin(caption="WARP", label="warp", standalone=False):
    if not standalone:
        print(r"""
  \caption{""" + caption + r"""}
  \label{fig:""" + label + r"""}
\end{figure}
""", end="")

def print_info(info):
    #print("% INFO " + info)
    pass


### PROPAGATION - TIKZ ###

def tikz_markbits(bits):
    print(r"""
  \begin{tikzpicture}[warpfig]
    \foreach \z in {""" + ",".join([str(bit) for bit in bits]) + r"""} { \fill[blue] (.25*\z,.75) circle[radius=3pt]; }
    \foreach \z[evaluate=\z as \zf using int(4*\z)] in {0,...,31} {
      \draw[gray] (\z,0) node[above] {\tiny\zf};
      \foreach \zb in {0,...,3} { \draw[gray] (\z+.25*\zb,0) -- +(0,-3pt); }
    }
  \end{tikzpicture}
""", end="")

def tikz_round(roundnr=0, annotations=""):
    print(r"""
  \warpround{""" + str(roundnr) + """}{""" + annotations + r"""}
""", end="")

def tikz_final(roundnr=0, annotations=""):
    print(r"""
  \warproundfinal{""" + str(roundnr) + """}{""" + annotations + r"""}
""", end="")


### PROPAGATION ###

def propagate_round(depends_in, perm=pi, tikz_fun=tikz_round, markoutput=False, roundnr=0):
    depends_out = []
    markbranches = []
    marksboxes = []
    for di in depends_in:
        depends_out.append(perm[di])
        markbranches.append((di,perm[di]))
        if di % 2:
            depends_out.append(perm[di-1])
            marksboxes.append((di-1,perm[di-1]))
    annotations = "%\n  "
    annotations += r"  \markbranches{" + ",".join([(str(di)+"/"+ str(pidi)) for (di,pidi) in markbranches]) + r"}" + "\n  "
    if marksboxes:
        annotations += r"  \marksboxes{" + ",".join([(str(di//2)+"/"+str(di)+"/"+ str(pidi)) for (di,pidi) in marksboxes]) + r"}" + "\n  "
    if markoutput:
        annotations += r"  \markoutputs{" + ",".join([str(pidi) for (di,pidi) in markbranches + marksboxes]) + r"}" + "\n  "
    tikz_fun(roundnr, annotations)
    return depends_out

def propagate_final(depends_in, roundnr=0):
    return propagate_round(depends_in, perm=list(range(32)), tikz_fun=tikz_final, markoutput=True, roundnr=roundnr)

def propagate_keyrecovery(nibbles=[], bits=[], rounds=7, distrounds=0, annotate=lambda: None):
    assert(bool(nibbles) ^ bool(bits))
    if nibbles:
        #nibbles = [int(nib) for nib in nibbles.split(",")]
        bits = sum([list(range(4*nib,4*nib+4)) for nib in nibbles], [])
    elif bits:
        #bits = [int(bit) for bit in bits.split(",")]
        nibbles = list(set([bit//4 for bit in bits]))

    tikz_markbits(bits)
    depends = [nibbles]

    for r in range(rounds-1):
        depends.append(propagate_round(depends[-1], roundnr=r+distrounds+1))
    depends.append(propagate_final(depends[-1], roundnr=rounds+distrounds))


### GRAPH CONSTRUCTION ###

class Cell:
    distrounds = 0

    def __init__(self, r, i, celltype="X"):
        self.r = r  # round index
        self.i = i  # cell index in round
        self.type = celltype  # X=intermediate, K=key, C=ciphertext
        self.t_label = self.type if self.type != "C" else "Y"
        self.r_label = r + distrounds if celltype != "K" else (r + distrounds) % 2
        self.i_label = i if celltype != "C" else pi.index(i) # hack: no permutation in last round
        self.label = "$\\" + self.t_label + "{" + str(self.r_label) + "}{" + str(self.i_label) + "}$"
        self.node = None
        self.synt = False


class SyntheticCell(Cell):
    def __init__(self, cell, sbox=False):
        self.r = -1
        self.i = -1
        self.type = cell.type
        self.label = cell.label if not sbox else r"$S(" + cell.label[1:-1] + r")$"
        self.cells = [cell]
        self.node = None
        self.synt = True

    def add_cell(self, cell, sbox=False):
        assert(self.type == cell.type)
        self.label = self.label[:-1] + "\\oplus" + (cell.label[1:] if not sbox else (r"S(" + cell.label[1:-1] + r")$"))
        self.cells.append(cell)
        

class Node:
    def __init__(self, cell, nodetype="tmp"):
        assert(cell.node is None)  # no two nodes for the same cell
        self.cell = cell  # corresponding state cell
        cell.node = self
        self.type = nodetype  # "tmp"=temporary intermediate, "xor"=xor of nodes, "key"=xor of keys, "out"=xor of ciphertext cells Ci or S(Ci)
        self.leaf = []
        self.stem = []

    def name(self):
        name = self.cell.label.strip("$")
        name = name.replace("}{", "i")
        for char in "\{} ()":
            name = name.replace(char, "")
        return name


    def tikz_node(self):
        dupflag = "dup," if len(self.stem) > 1 else ""
        if self.type == "xor":
            return r"node[" + dupflag + r"XOR,label=" + self.cell.label + r"] (" + self.name() + ") {+}"
        elif self.type == "out":
            return r"node[" + dupflag + r"OUT] (" + self.name() + ") {" + self.cell.label + r"}"
        elif self.type == "key":
            return r"node[" + dupflag + r"KEY] (" + self.name() + ") {" + self.cell.label + r"}"
        else:
            return r"node[" + dupflag + r"] {$\bullet$}"
        
    def tikz_recursive(self, indent, end=""):
        print(self.tikz_node(), end="")
        for child in self.leaf:
            print("\n" + indent + r"child {", end="")
            child.leaf.tikz_recursive(indent + "       ", end=end)
            if child.type == "S":
                print("\n" + indent + r"       edge from parent node[Sbox] {$S$}", end=end)
            print(r"}", end=end)


    def tikz_graph_node(self, opt=""):
        flag = (opt + ",") if opt else ""
        flag += ("dup," if len(self.stem) > 1 else "")
        if self.type == "xor":
            return self.name() + "/+"                  + r" [" + flag + r"XOR,label=" + self.cell.label + r"]"
        elif self.type == "out":
            return self.name() + "/" + self.cell.label + r" [" + flag + r"OUT]"
        elif self.type == "key":
            return self.name() + "/" + self.cell.label + r" [" + flag + r"KEY]"
        else:
            return r"$\bullet$ [" + dupflag + r"]"
        
    def tikz_graph_recursive(self, indent, opt="", end=""):
        print(indent + self.tikz_graph_node(opt), end="")
        if self.leaf:
            print(r" -> {")
            for child in self.leaf:
                if child.type == "S":
                    child.leaf.tikz_graph_recursive(indent + "         ", opt="SBX", end=",\n")
                    #print("\n" + indent + r"       edge from parent node[Sbox] {$S$}", end=end)
                else:
                    child.leaf.tikz_graph_recursive(indent + "         ", end=",\n")
            print(indent + r"}", end=end)
        else:
            print("", end=end)


class Edge:
    def __init__(self, stem, leaf, optype="I"):
        self.type = optype  # "S"=Sbox, "I"=identity
        self.stem = stem  # parent node, computed from edge
        self.leaf = leaf  # child node, input for computation
        leaf.stem.append(self)
        stem.leaf.append(self)


class Graph:
    def __init__(self, root):
        self.root = root
        self.node = [root]
        self.edge = []

    def add_node(self, cell, stem, sbox=False):
        """create Node for cell if necessary, set its stem, and return node for this cell"""
        add_edge = True
        if cell.type in ["K", "C"] and [sib for sib in stem.leaf if sib.leaf.cell.type == cell.type]:
            # if parent already has a child of same type 'out' or 'key' -> synthetic cell
            sibs = [sib for sib in stem.leaf if sib.leaf.cell.type == cell.type]
            assert(len(sibs) == 1)
            sib = sibs[0].leaf
            if sib.cell.synt: 
                assert(len(sib.stem) == 1) # need to expand code if this happens!
                sib.cell.add_cell(cell, sbox)
                node = sib
                print_info("B sib updated: " + node.name() + " with " + cell.label)
                add_edge = False
            else:
                syntcell = SyntheticCell(sib.cell, sbox=(sib.stem[0].type == "S")) # FIX: not stem[0] but the right stem!
                syntcell.add_cell(cell, sbox)
                samecell = [node for node in self.node if node.cell.label == syntcell.label] # this is not robust under permutations, would need some kind of sorting/unique representation - doesn't appear in our examples
                if samecell:
                    # different node with same synthetic cells already exists -> use node, remove sib
                    node = samecell[0]
                    stem.leaf = [oldleaf for oldleaf in stem.leaf if oldleaf.leaf != sib]
                    sib.stem = [oldstem for oldstem in sib.stem if oldstem.stem != stem]
                    sbox = False  # edge will be "I" edge in any case
                    print_info("C samecell: " + node.name())
                    if not sib.stem:
                        # sibling node not needed anymore
                        print_info("C and remove sib: " + sib.name())
                        sib.cell.node = None
                        self.node.remove(sib)
                elif len(sib.stem) == 1:
                    # switch cell in existing node
                    sib.cell.node = None
                    syntcell.node = sib
                    sib.cell = syntcell
                    sib.stem[0].type = "I"
                    node = sib
                    add_edge = False
                    print_info("D switch cell of node: " + node.name())
                else:
                    # create new node, remove old edge
                    node = Node(syntcell, nodetype=sib.type)
                    self.node.append(node)
                    stem.leaf = [oldleaf for oldleaf in stem.leaf if oldleaf.leaf != sib]
                    sib.stem = [oldstem for oldstem in sib.stem if oldstem.stem != stem]
                    print_info("E create new node: " + node.name())
                    sbox = False  # edge will be "I" edge in any case
        elif cell.node: # node exists already
            node = cell.node
            print_info("A exists: " + node.name())
        else: # create new node
            nodetype = {"X": "tmp", "K": "key", "C": "out"}
            node = Node(cell, nodetype=nodetype[cell.type])
            self.node.append(node)
            print_info("F create new node: " + node.name())
        if add_edge:
            edge = Edge(stem, node, "S" if sbox else "I")
            self.edge.append(edge)
        return node

    def tikz(self, annotate=lambda: None):
        print(r"""
  \begin{tikzpicture}[tree layout,
                      grow=right,
                      sibling distance=1cm, level distance=2cm,
                      every label/.append style={gray,font=\footnotesize,below left},
                      Sbox/.style={circle,draw=black,fill=white,inner sep=1pt},
                      dup/.style={draw=red,dashed,thick},
                      KEY/.style={blue},
                      OUT/.style={green!50!black},
                      XOR/.style={circle,draw,inner sep=0pt}]
    """ + "\\", end="")
        self.root.tikz_recursive("     ")
        print(r"""
     ;
    \draw node[right,align=left] {%""", end="")
        annotate()
        print(r"""};
  \end{tikzpicture}
""", end="")

    def tikz_graph(self, annotate=lambda: None):
        print(r"""
  \begin{tikzpicture}[tree layout,
                      grow=right,
                      >=latex,
                      sibling distance=.5cm, level distance=2cm,
                      every edge quotes/.style={fill=white,draw,circle,inner sep=1pt},
                      every label/.append style={gray,font=\footnotesize,below left},
                      SBX/.style={>"S"},
                      dup/.style={draw=red,dashed,thick},
                      KEY/.style={blue},
                      OUT/.style={green!50!black},
                      XOR/.style={circle,draw,inner sep=0pt}]
    """ + "\\graph {")
        self.root.tikz_graph_recursive("    ")
        print(r"""
    };
    \draw node[right,align=left] {%""", end="")
        annotate()
        print(r"""};
  \end{tikzpicture}
""", end="")


### COMPLEXITY ###

class KeyRecovery:
    statesize = 32

    def __init__(self, rounds, perm=pi, distrounds=22, data_log=127):
        self.rounds = rounds
        self.perm = perm
        self.distrounds = distrounds
        Cell.distrounds = distrounds
        self.cells = [[Cell(r,i,"X") for i in range(32)] for r in range(rounds)] + [[Cell(rounds,i,"C") for i in range(32)]]  # state cells
        self.keys = [[Cell(r,i,"K") for i in range(16)] for r in range(2)]  # key cells
        self.extra = []  # synthetic cells for key recovery, e.g., xor of keys
        self.data_log = data_log
        self.graph = None

    def recover_cell(self, i, r=0, merge=True):
        """determine the required cells to recover cell i in the output of round r (default: before the first round)"""
        root = Node(self.cells[r][i])
        self.graph = Graph(root)
        tmp_nodes = [root]
        while tmp_nodes:
            tmp_nodes = sum([self.expand_node(node, merge) for node in tmp_nodes if node.type == "tmp"], [])

    def expand_node(self, node, merge=True):
        """expand a new node's dependencies and return list of any new nodes it depends on"""
        if node.cell.i % 2 == 0 and node.cell.type == "X":  # is even-indexed identity node -> expand
            node.cell = self.cells[node.cell.r+1][self.perm[node.cell.i]]
            node.cell.node = node
            if node.cell.type == "C":
                node.type = "out"
        if node.type == "tmp":  # new node -> expand
            assert(node.cell.type == "X" and node.cell.i % 2)
            stem = node
            cell = node.cell
            if len(node.stem) == 1 and node.stem[0].stem.type == "xor" and node.stem[0].type == "I":
                if merge:
                    # merge xor nodes
                    stem = node.stem[0].stem
                    self.graph.node.remove(node)
                    self.graph.edge.remove(node.stem[0])
                    stem.leaf.remove(node.stem[0])
                    cell.node = None
                    print_info("M merge xor node " + node.name() + " into parent " + stem.name())

            node.type = "xor"  # mark as resolved (even if deleted)
            node_L = self.graph.add_node(self.cells[cell.r+1][self.perm[cell.i-1]], stem, sbox=True)
            node_R = self.graph.add_node(self.cells[cell.r+1][self.perm[cell.i]], stem)
            node_K = self.graph.add_node(self.keys[cell.r % 2][cell.i // 2], stem)
            print_info("X node " + node.name() + " adds " + node_L.name() + ", " + node_R.name() + ", " + node_K.name())
            return [node_L, node_R, node_K]
        return []


    def get_equation(self, cell):
        if cell.type == "C":
            i_one = [cell.i] if not cell.synt else [ci.i for ci in cell.cells]
        elif cell.type == "K":
            i_one = [16*cell.r + cell.i] if not cell.synt else [16*ci.r+ci.i for ci in cell.cells]
        else:
            assert(cell.type in "CK")
        return [(1 if i in i_one else 0) for i in range(32)]

    def get_num_C(self):
        """number of ciphertext nibbles involved in key recovery"""
        C_cells = [node for node in self.graph.node if node.type == "out"]
        for i, C_cell in enumerate(C_cells):
            print_info("C nibble " + str(i) + ": " + C_cell.cell.label)
        if WITH_SAGE:
            #print_info("Key matrix:\n" + str(Matrix(GF(2), [self.get_equation(Ci.cell) for Ci in C_cells])))
            return Matrix(GF(2), [self.get_equation(Ci.cell) for Ci in C_cells]).rank()
        else:
            return len(set(C_cells))
            #return len(C_cells)

    def get_num_K(self):
        """number of key nibbles involved in key recovery"""
        #K_cells = [node for node in self.graph.node if node.type == "key" and not (node.stem[0].stem == self.graph.root and node.stem[0].type == "I")]   # - excludes linear key, not applicable for X_L
        K_cells = [node for node in self.graph.node if node.type == "key"]
        for i, K_cell in enumerate(K_cells):
            print_info("K nibble " + str(i) + ": " + K_cell.cell.label)
        if WITH_SAGE:
            #print_info("Key matrix:\n" + str(Matrix(GF(2), [self.get_equation(Ki.cell) for Ki in K_cells])))
            return Matrix(GF(2), [self.get_equation(Ki.cell) for Ki in K_cells]).rank()
        else:
            return len(set(K_cells))
            #return len(K_cells)

    def get_num_S(self):
        """number of Sbox calls involved in partial decryption"""
        S_edges = [edge for edge in self.graph.edge if edge.type == "S"]
        return len(S_edges)

    def get_S_ratio(self):
        """encryption call equivalents of one partial decryption"""
        return self.get_num_S() / ((self.rounds+self.distrounds) * 16)


    def complexity_plain(self, print_tex=True):
        """group ciphertexts by their relevant output bits; for each key candidate for each ciphertext value, compute back and check sum"""
        bit_C = 4*self.get_num_C()
        bit_K = 4*self.get_num_K()
        ratio = self.get_S_ratio()
        if print_tex:
            print(r"""
\paragraph{Attack complexity (plain)}~\\
Data: $2^{""" + str(self.data_log) + r"""}$ queries\\
Time: $""" + str(ratio) + r""" \cdot 2^{""" + str(bit_K) + r"+" + str(bit_C) + r"""} = 2^{""" + str(log(ratio,2) + bit_K + bit_C) + r"""}$ decryptions (plus queries)\\
Mem:  $2^{""" + str(bit_C) + r"""}$ partial ciphertexts
""", end="")
        return bit_C + bit_K

    def complexity_fft(self, print_tex=True):
        """use FFT for integral attacks"""
        bit_C = 4*self.get_num_C()
        bit_K = 4*self.get_num_K()
        if print_tex:
            print(r"""
Data: $2^{""" + str(self.data_log) + r"""}$ queries\\
Keys: $2^{""" + str(bit_K) + r"""}$ key candidates\\
Mem:  $2^{""" + str(bit_C) + r"""}$ partial ciphertexts\\
Time: $2^{""" + str(bit_K+log(bit_K,2)) + r"""}$ additions (ideally)
""", end="")
        return max(bit_C, log(bit_K,2) + bit_K)


    def print_cipher(self, standalone=False, complexity=lambda: None):
        assert(self.graph)
        tex_figure_init(standalone=standalone)
        propagate_keyrecovery(nibbles=[self.graph.root.cell.i], rounds=self.rounds, distrounds=self.distrounds, annotate=complexity)
        tex_figure_fin(r"Key recovery for " + str(rounds) + r" rounds of \texttt{WARP}.", r"warp_propagation", standalone=standalone)

    def print_graph(self, standalone=False, complexity=lambda: None):
        assert(self.graph)
        tex_figure_init(standalone=standalone)
        #self.graph.tikz(annotate=complexity)
        self.graph.tikz_graph(annotate=complexity)
        tex_figure_fin(r"Key guessing DAG for " + str(self.rounds) + r" rounds", r"warp_dag_" + str(self.rounds), standalone=standalone)


### MAIN ###

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-R", "--rounds", default=9, type=int, help="The number of key-recovery rounds")
    parser.add_argument("-r", "--distrounds", default=23, type=int, help="The number of distinguisher rounds")
    parser.add_argument("-d", "--data", default=127, type=int, help="log2(data complexity) of distinguisher")

    target_group = parser.add_mutually_exclusive_group()
    target_group.add_argument("-n", "--nibble", default=5, type=int, help="The target zero-sum nibble (0..31)")
    target_group.add_argument("-b", "--bit", default=None, type=int, help="The target zero-sum bit (0..127)")

    parser.add_argument("-m", "--method", default="fft", help="Key recovery method ('naive', 'fft') - currently ignored")
    parser.add_argument("-N", "--nomerge", action="store_true", help="don't merge XOR nodes")
    parser.add_argument("-c", "--printcipher", action="store_true", help="print cipher figure")
    parser.add_argument("-g", "--printgraph", action="store_true", help="print key recovery DAG")
    parser.add_argument("-s", "--standalone", action="store_true", help="standalone document type")
    return vars(parser.parse_args())

if __name__ == "__main__":
    locals().update(parse_args())

    tex_init(" ".join(sys.argv), standalone=standalone)
    keyrec = KeyRecovery(rounds=rounds, distrounds=distrounds, data_log=127)
    keyrec.recover_cell(nibble, merge=not nomerge)
    complexities = {"naive": keyrec.complexity_plain,
                    "fft": keyrec.complexity_fft}
    if printcipher:
        keyrec.print_cipher(standalone=standalone, complexity=complexities[method])
    if printgraph:
        keyrec.print_graph(standalone=standalone, complexity=complexities[method])
    if not printcipher and not printgraph:
        complexities[method]()

    tex_fin(rounds)