#!/usr/bin/python

# author: Quentin L. Meunier

from __future__ import print_function
import re
import os
import sys
import subprocess


nb_bits = 3

indent_base = 4

def write_to_file(content, filename):
    print("# Writing to file %s" % (filename))
    f = open(filename, "w")
    f.write(content)
    f.close()


def print_and_call(cmd):
    print(subprocess.list2cmdline(cmd))
    retval = subprocess.call(cmd)
    return retval

def print_and_popen(cmd, outfile):
    print(subprocess.list2cmdline(cmd))
    print(" > " + outfile)
    output = subprocess.Popen(cmd, stdout = subprocess.PIPE).communicate()[0]
    return output


def indent(level, s):
    return level * indent_base * ' ' + s



def header():
    ret = '''
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>


int32_t sra(int32_t v, int32_t sh) {
    return ((v << %(sh)d) >> %(sh)d) >> sh;
}


int32_t main() {
''' % dict(sh = 32 - nb_bits)
    return ret


def init_distrib_tab(formula, nb_bits, secrets):
    max_val = (1 << nb_bits)
    indent_level = 1
    ret = indent(indent_level, "// Distribution of expression: %s\n" % (formula))
    ret += indent(indent_level, "{\n")
    indent_level += 1
    ret += indent(indent_level, "int32_t distrib[%d]" % (max_val))
    for s in secrets:
        ret += "[%d]" % max_val
    ret += ";\n"
    for s in secrets:
        ret += indent(indent_level, "for (int32_t %s = 0; %s < %d; %s++) {\n" % (s, s, max_val, s))
        indent_level += 1
    ret += indent(indent_level, "for (int32_t i = 0; i < %d; i++) {\n" % (max_val))
    indent_level += 1
    ret += indent(indent_level, "distrib")
    for s in secrets:
        ret += "[%s]" % s
    ret += "[i] = 0;\n"
    for s in secrets:
        indent_level -= 1
        ret += close_curlybrace(indent_level)
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    return ret

    
def formula_computation(formula, nb_bits, indent_level):
    return indent(indent_level, "int32_t res = (%s) & 0x%x;\n" % (formula, (1 << nb_bits) - 1))


def check_distrib(secrets, masks, formula, nb_bits):
    nb_masks = len(masks)
    max_val = (1 << nb_bits)
    indent_level = 2
    ret = indent(indent_level, "int32_t nb_diff = 0;\n");
    ret += indent(indent_level, "int32_t uniform = 1;\n")
    for s in secrets:
        ret += indent(indent_level, "for (int32_t %s = 0; %s < %d; %s++) {\n" % (s, s, max_val, s))
        indent_level += 1

    ret += indent(indent_level, "// Checking distribution for each key combination\n")
    ret += indent(indent_level, "for (int32_t val = 0; val < %d; val++) {\n" % (max_val))
    indent_level += 1
    ret += indent(indent_level, "if (distrib")
    for s in secrets:
        ret += "[%s]" % s
    ret += "[val] != %d) {\n" % ((1 << nb_bits) ** (nb_masks - 1))
    indent_level += 1
    ret += indent(indent_level, "uniform = 0;\n")
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    ret += indent(indent_level, "if (distrib")
    for s in secrets:
        ret += "[%s]" % s
    ret += "[val] != distrib" + "[0]" * len(secrets) + "[val]) {\n"
    indent_level += 1
    ret += indent(indent_level, "nb_diff += 1;\n")
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    for s in secrets:
        indent_level -= 1
        ret += close_curlybrace(indent_level)

    ret += indent(indent_level, "if (uniform == 1) {\n");
    indent_level += 1
    ret += indent(indent_level, "printf(\"[RUD] %s\\n\");\n" % formula)
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    ret += indent(indent_level, "else if (nb_diff == 0) {\n")
    indent_level += 1
    ret += indent(indent_level, "printf(\"[ISD] %s\\n\");\n" % formula)
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    ret += indent(indent_level, "else {\n")
    indent_level += 1
    ret += indent(indent_level, "printf(\"[UKD] %s\\n\");\n" % formula)
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    indent_level -= 1
    ret += close_curlybrace(indent_level)
    return ret


def for_loop(v, max_val, indent_level):
    return indent(indent_level, "for (int32_t %s = 0; %s < %d; %s++) {\n" % (v, v, max_val, v))

def close_curlybrace(indent_level):
    return indent(indent_level, "}\n")

def incr_distrib(secrets, ident_level):
    ret = indent(indent_level, "distrib")
    for s in secrets:
        ret += "[%s]" % s
    ret += "[res] += 1;\n"
    return ret

def usage(name):
    print("usage: %s [-f <filename> | -l <formula> ]" % name)
    print("Check distribution of one or several formula(s) using an exhaustive enumeration.")
    print("Options:")
    print("   -h            : display this help")
    print("   -f <filename> : filename is a file which contains the formulas to check")
    print("   -l <formula>  : formula to check directly on the command line")
    print("Names of secret variables must be in the format k[0-9]* or s[0-9]* and names of masks must in the format m[0-9]* of r[0-9]*")
    print("Accepted operators are: +, -, &, |, ^, ~, (, ), <<, >>, sra(...)")

def parse_args(argv):
    i = 1
    if len(argv) == 1:
        usage(argv[0])
        sys.exit(1)
    while i != len(argv):
        #print("analyzing arg %s" % argv[i])
        if argv[i] == "-h":
            usage(argv[0])
            sys.exit(0)
        elif argv[i] == "-f":
            i += 1
            c_file = argv[i] + ".c"
            exec_file = argv[i] + ".x"
            # check and bla
            with open(argv[i]) as f:
                formulas = f.read().splitlines()
        elif argv[i] == "-l":
            i += 1
            formulas = list()
            formulas.append(argv[i])
            c_file = "a.c"
            exec_file = "a.x"
        else:
            print("*** Error: unknown option: %s", argv[i])
            usage(argv[0])
            sys.exit(1)
        i += 1
    return formulas, c_file, exec_file


formulas, c_file, exec_file = parse_args(sys.argv)

if len(formulas) == 0:
    print("*** Error: no file or formula given")
    usage(sys.argv)
    sys.exit(1)

content = header()

for formula in formulas:
    #print("formula : %s" % formula)
    if formula == "":
        continue
    exps = set() # non terminal expressions
    masks = set()
    secrets = set()
    exps.add(formula)
    separators = [' ', 'sra', '+', '-', '&', '|', '^', '~', '(', ')', '<<', '>>', ',']
    for separator in separators:
        new_exps = set()
        for exp in exps:
            results = exp.split(separator)
            for result in results:
                if result == '':
                    continue
                terminal = re.search('^[ks][0-9]*$', result)
                if terminal != None:
                    secrets.add(result)
                else:
                    terminal = re.search('^[mr][0-9]*$', result)
                    if terminal != None:
                        masks.add(result)
                    # Not variable name, testing for constants
                    else:
                        terminal = re.search('^(|0x)[0-9a-f]+$', result)
                        if terminal == None:
                            # Not a constant, expression needs more splitting
                            new_exps.add(result)
        exps = new_exps
    if len(exps) != 0:
        print("Errors:")
        for exp in exps:
            print("expression %s cannot be processed" % exp)
            sys.exit(1)
    
    content += init_distrib_tab(formula, nb_bits, secrets)

    indent_level = 2
    for variable in secrets:
        content += for_loop(variable, (1 << nb_bits), indent_level)
        indent_level += 1
    for variable in masks:
        content += for_loop(variable, (1 << nb_bits), indent_level)
        indent_level += 1
        
    content += formula_computation(formula, nb_bits, indent_level)
    #content += ' ' * (indent_level * 4) + 'printf("val : %d\\n", res);\n'
    content += incr_distrib(secrets, indent_level)
    
    for variable in secrets:
        indent_level -= 1
        content += close_curlybrace(indent_level)
    for variable in masks:
        indent_level -= 1
        content += close_curlybrace(indent_level)

    content += check_distrib(secrets, masks, formula, nb_bits) 


indent_level -= 1
content += indent(indent_level, "return 0;\n")
content += "}\n\n"

write_to_file(content, c_file)

cmd = ['gcc', '-O2', '-std=c99', '-o', exec_file , c_file]
retval = print_and_call(cmd)

if retval != 0:
    print("*** Error while executing command %s" %cmd)
    sys.exit(1)

cmd = ['./' + exec_file]
print_and_call(cmd)