import math
import sys
from sets import Set
import xlwt


Sboxes = [ [0x3E, 0x72, 0x5B, 0x47, 0xCA, 0xE0, 0x00, 0x33, 0x04, 0xD1, 0x54, 0x98, 0x09, 0xB9, 0x6D, 0xCB,
            0x7B, 0x1B, 0xF9, 0x32, 0xAF, 0x9D, 0x6A, 0xA5, 0xB8, 0x2D, 0xFC, 0x1D, 0x08, 0x53, 0x03, 0x90,
            0x4D, 0x4E, 0x84, 0x99, 0xE4, 0xCE, 0xD9, 0x91, 0xDD, 0xB6, 0x85, 0x48, 0x8B, 0x29, 0x6E, 0xAC,
            0xCD, 0xC1, 0xF8, 0x1E, 0x73, 0x43, 0x69, 0xC6, 0xB5, 0xBD, 0xFD, 0x39, 0x63, 0x20, 0xD4, 0x38,
            0x76, 0x7D, 0xB2, 0xA7, 0xCF, 0xED, 0x57, 0xC5, 0xF3, 0x2C, 0xBB, 0x14, 0x21, 0x06, 0x55, 0x9B,
            0xE3, 0xEF, 0x5E, 0x31, 0x4F, 0x7F, 0x5A, 0xA4, 0x0D, 0x82, 0x51, 0x49, 0x5F, 0xBA, 0x58, 0x1C,
            0x4A, 0x16, 0xD5, 0x17, 0xA8, 0x92, 0x24, 0x1F, 0x8C, 0xFF, 0xD8, 0xAE, 0x2E, 0x01, 0xD3, 0xAD,
            0x3B, 0x4B, 0xDA, 0x46, 0xEB, 0xC9, 0xDE, 0x9A, 0x8F, 0x87, 0xD7, 0x3A, 0x80, 0x6F, 0x2F, 0xC8,
            0xB1, 0xB4, 0x37, 0xF7, 0x0A, 0x22, 0x13, 0x28, 0x7C, 0xCC, 0x3C, 0x89, 0xC7, 0xC3, 0x96, 0x56,
            0x07, 0xBF, 0x7E, 0xF0, 0x0B, 0x2B, 0x97, 0x52, 0x35, 0x41, 0x79, 0x61, 0xA6, 0x4C, 0x10, 0xFE,
            0xBC, 0x26, 0x95, 0x88, 0x8A, 0xB0, 0xA3, 0xFB, 0xC0, 0x18, 0x94, 0xF2, 0xE1, 0xE5, 0xE9, 0x5D,
            0xD0, 0xDC, 0x11, 0x66, 0x64, 0x5C, 0xEC, 0x59, 0x42, 0x75, 0x12, 0xF5, 0x74, 0x9C, 0xAA, 0x23,
            0x0E, 0x86, 0xAB, 0xBE, 0x2A, 0x02, 0xE7, 0x67, 0xE6, 0x44, 0xA2, 0x6C, 0xC2, 0x93, 0x9F, 0xF1,
            0xF6, 0xFA, 0x36, 0xD2, 0x50, 0x68, 0x9E, 0x62, 0x71, 0x15, 0x3D, 0xD6, 0x40, 0xC4, 0xE2, 0x0F,
            0x8E, 0x83, 0x77, 0x6B, 0x25, 0x05, 0x3F, 0x0C, 0x30, 0xEA, 0x70, 0xB7, 0xA1, 0xE8, 0xA9, 0x65,
            0x8D, 0x27, 0x1A, 0xDB, 0x81, 0xB3, 0xA0, 0xF4, 0x45, 0x7A, 0x19, 0xDF, 0xEE, 0x78, 0x34, 0x60 ],
            [0x55, 0xC2, 0x63, 0x71, 0x3B, 0xC8, 0x47, 0x86, 0x9F, 0x3C, 0xDA, 0x5B, 0x29, 0xAA, 0xFD, 0x77,
            0x8C, 0xC5, 0x94, 0x0C, 0xA6, 0x1A, 0x13, 0x00, 0xE3, 0xA8, 0x16, 0x72, 0x40, 0xF9, 0xF8, 0x42,
            0x44, 0x26, 0x68, 0x96, 0x81, 0xD9, 0x45, 0x3E, 0x10, 0x76, 0xC6, 0xA7, 0x8B, 0x39, 0x43, 0xE1,
            0x3A, 0xB5, 0x56, 0x2A, 0xC0, 0x6D, 0xB3, 0x05, 0x22, 0x66, 0xBF, 0xDC, 0x0B, 0xFA, 0x62, 0x48,
            0xDD, 0x20, 0x11, 0x06, 0x36, 0xC9, 0xC1, 0xCF, 0xF6, 0x27, 0x52, 0xBB, 0x69, 0xF5, 0xD4, 0x87,
            0x7F, 0x84, 0x4C, 0xD2, 0x9C, 0x57, 0xA4, 0xBC, 0x4F, 0x9A, 0xDF, 0xFE, 0xD6, 0x8D, 0x7A, 0xEB,
            0x2B, 0x53, 0xD8, 0x5C, 0xA1, 0x14, 0x17, 0xFB, 0x23, 0xD5, 0x7D, 0x30, 0x67, 0x73, 0x08, 0x09,
            0xEE, 0xB7, 0x70, 0x3F, 0x61, 0xB2, 0x19, 0x8E, 0x4E, 0xE5, 0x4B, 0x93, 0x8F, 0x5D, 0xDB, 0xA9,
            0xAD, 0xF1, 0xAE, 0x2E, 0xCB, 0x0D, 0xFC, 0xF4, 0x2D, 0x46, 0x6E, 0x1D, 0x97, 0xE8, 0xD1, 0xE9,
            0x4D, 0x37, 0xA5, 0x75, 0x5E, 0x83, 0x9E, 0xAB, 0x82, 0x9D, 0xB9, 0x1C, 0xE0, 0xCD, 0x49, 0x89,
            0x01, 0xB6, 0xBD, 0x58, 0x24, 0xA2, 0x5F, 0x38, 0x78, 0x99, 0x15, 0x90, 0x50, 0xB8, 0x95, 0xE4,
            0xD0, 0x91, 0xC7, 0xCE, 0xED, 0x0F, 0xB4, 0x6F, 0xA0, 0xCC, 0xF0, 0x02, 0x4A, 0x79, 0xC3, 0xDE,
            0xA3, 0xEF, 0xEA, 0x51, 0xE6, 0x6B, 0x18, 0xEC, 0x1B, 0x2C, 0x80, 0xF7, 0x74, 0xE7, 0xFF, 0x21,
            0x5A, 0x6A, 0x54, 0x1E, 0x41, 0x31, 0x92, 0x35, 0xC4, 0x33, 0x07, 0x0A, 0xBA, 0x7E, 0x0E, 0x34,
            0x88, 0xB1, 0x98, 0x7C, 0xF3, 0x3D, 0x60, 0x6C, 0x7B, 0xCA, 0xD3, 0x1F, 0x32, 0x65, 0x04, 0x28,
            0x64, 0xBE, 0x85, 0x9B, 0x2F, 0x59, 0x8A, 0xD7, 0xB0, 0x25, 0xAC, 0xAF, 0x12, 0x03, 0xE2, 0xF2 ] ]


def createEmptyList(size):
    return [0]*size


def createEmpty2DList(size):
    result = []
    x = 0
    while x < size:
        result.append(createEmptyList(size))
        x += 1
    return result


def generatePairsWithDifference( diff ):
    result = []
    for x in range (0,256):
        for y in range (0,256):
            if x ^ y == diff:
               result.append((x,y))
               break 
    return result


def outputDifferenceOfPair(pair, SBoxes, SBoxNumber):
    firstOutput = SBoxes[SBoxNumber][pair[0]]
    secondOutput = SBoxes[SBoxNumber][pair[1]] 
    return firstOutput ^ secondOutput


#########################################################################################################
def printDDT( DDT , sheet):
    # print "\t|0\t1\t2\t3\t4\t5\t6\t7\t8\t9\tA\tB\tC\tD\tE\tF"
    print " |",
    for i in range(0,256):
        print "%d|"%i,

    print "----------------------------------------------------------------------------------------------------------------------------------"
    
    k = 0
    while k < len(DDT):
        s = str(k) + "|"
        col = 0
        for i in DDT[k]:
            s = s + "" + str(i) + "|"
            sheet.write(k, col, str(i))
            col = col + 1
        print s
        k = k + 1
    
    return sheet


def DDT(SBoxes, SBoxNumber):
    DDT = createEmpty2DList(256)
    differentialUniformity = 0   
    for x in range(0,256):
        pairs = generatePairsWithDifference(x)
        for pair in pairs:
            diff = outputDifferenceOfPair(pair, SBoxes, SBoxNumber)
            DDT[x][diff] += 1
            if x != 0 and DDT[x][diff] > differentialUniformity:
                differentialUniformity = DDT[x][diff]
    print "\nDifferential Uniformity of SBox" + str(SBoxNumber) + " : " + str(differentialUniformity)
    return DDT


def printAllDDT( SBoxes, typ ):
    # create excel
    wbk = xlwt.Workbook()
    for i in range(0,2):
        print "#### Differential Distribution Table of " + typ + "SBox-" + str(i) + " ####"        
        # create 2 DDT excel table's tittle
        sheet = wbk.add_sheet('sheet %d' %(i+1))            
        sheet = printDDT(DDT(SBoxes, i), sheet)
        print "\n"
    wbk.save('DDT.xls')

#########################################################################################################
        
        
def printLAT( LAT ):
    print " |",
    for i in range(0,256):
        print "%d|"%i,
    
    print "----------------------------------------------------------------------------------------------------------------------------------"
    i = 0
    line = ""
    for entry in LAT:
        if i%256 == 0:
            line += str(i/256) + "|"
        line += str(entry) + "|"
        if i%256 == 255:
            line += "\n"
        i += 1
    print line
        
def mask(x, a):
    binX = bin(x)[2:]
    binA = bin(a)[2:]
    
    if len(binX) < 8:
        binX = "0"*(8-len(binX)) + binX
    
    if len(binA) < 8:
        binA = "0"*(8-len(binA)) + binA
    
    r = int(binX[0]) & int(binA[0])
    for j in range(1,8):
        r ^= int(binX[j]) & int(binA[j])
    return r
        
def LATofSBOX( sBox ,sheet):   
    lat = []
    for i in range(0,256):
        for j in range(0,256):
            r = 0
            for x in range(0,256):
                sx = sBox[x]
                r += (mask(i,x) ^ mask(j,sx))
            r = -(r - 128)
            lat.append(r)
            
            sheet.write(i, j, r)
        print i
    printLAT(lat)
    return sheet

def LAT(sBoxes):
    # create excel
    wbk = xlwt.Workbook()
    
    i = 0
    while i < len(sBoxes):
        print "######## LAT of SBox-" + str(i) + " ########"
        
        # create 2 LAT excel table's tittle
        sheet = wbk.add_sheet('sheet %d' %(i+1))
        
        sheet = LATofSBOX( sBoxes[i] ,sheet)
        i += 1
    
    wbk.save('LAT.xls')

################################# MAIN #############################################3


if sys.argv[1] == "DDT" : 
    printAllDDT(Sboxes, "")
elif sys.argv[1] == "LAT":
    LAT( Sboxes )