Protocolo DH

from sys import argv

def potmod(x, pot, mod):
    bits = pot
    res = 1
    temp = x
    while bits > 0:
        if bits %2 == 1:
            res = (res * temp  % mod)
        temp = (temp * temp) % mod
        bits = bits >> 1
    return res

def bruteForce(generador, primo, F):
    for i in range(primo):
        valor = potmod(generador, i, primo)
        if valor == F:
            return i

def imprimeBruteForce(generador, primo, F):
    resultado = ''
    for i in range(primo):
        valor = potmod(generador, i, primo)
        if valor == F:
            resultado = 'El numero utilizado es:'+str(i)
        print '%d^%d mod %d = %d' % (generador, i, primo, valor)
        print resultado
        print 'Valor de Funcion buscado: %d' % F
    return

def testGenerador(z, p):
    lista = []
    for i in range(1,p):
        lista.append(z*i % p)
    for j in range(1,p):
        if lista.count(j) > 1:
            return False
    return True

if __name__ == "__main__":
    potA = int(argv[1])
    potB = int(argv[2])
    generador = int(argv[3])
    primo = int(argv[4])

    X = potmod(generador, potA, primo)
    Ka = X**potB % primo

    Y = potmod(generador, potB, primo)
    Kb = Y**potA % primo

    print '%d mod %d primo es gnerador:' %(generador, primo),
    print testGenerador(generador, primo)

    if testGenerador(generador, primo):

        print 'El valor de la funcion para Alice es %d y el valor de K es %d' % (X,Ka)
        print 'El valor de la funcion para Bob es %d y el valor de K es %d' % (Y,Kb)

        print 'La potencia para Alice original es %d y la potencia calculada es de %d' % (potA, bruteForce(generador, primo, X))
        print 'La potencia para Bob original es %d y la potencia calculada es de %d' % (potB, bruteForce(generador, primo, Y))