
# DSA.py : Digital Signature Algorithm implementation
#
# Distributed under the terms described in the LICENCE file.
#
# Written by akuchlin@cnri.reston.va.us.
#
# Implements a completely non-standards-compliant version of the
# Digital Signature Algorithm, which is used to sign
# the authentication tokens handed out by the central Web server at
# CNRI.
#
# This could easily be made configurable, in order to use Kerberos, properly
# licensed RSA, or whatever other authenication techniques are desired.
# (I'm tempted to drop DSA and just use a keyed hash to generate the
# authentication tokens, but will stick with DSA for this release.

import types

error = 'DSA error'

# Two functions to convert between long integers and strings

def long2str(n):
    s=''
    while n>0:
        s=chr(n & 255)+s
        n=n>>8
    return s

def str2long(s):
    if type(s)!=types.StringType: return s   # Integers will be left alone
    return reduce(lambda x,y : x*256+ord(y), s, 0L)

# Various number-theoretical utility functions to compute inverses,
# check primality, and so forth.  Most are the same algorithms
# described in _Applied Cryptography_.

def inverse(u, v):
    "Return the inverse of u mod v."
    u3, v3 = long(u), long(v)
    u1, v1 = 1L, 0L
    while v3>0:
	q=u3/v3
	u1, v1 = v1, u1-v1*q
	u3, v3 = v3, u3-v3*q
    while u1<0: u1=u1+v
    return u1

def getRandomNumber(N, randfunc):
    "Return an N-bit random number."
    str=randfunc(N/8)
    char=ord(randfunc(1))>>(8-(N%8))
    return str2long(chr(char)+str)
    
# Given a number of bits to generate and a random generation function,
# find a prime number of the appropriate size.

def getPrime(N, randfunc):
    "Return a random N-bit prime number."
    number=getRandomNumber(N, randfunc) | 1
    while (not isPrime(number)):
        number=number+2
    return number

def isPrime(N):
    "Return true if N is prime."
    for i in sieve:
        if (N % i)==0: return 0

    # Compute the highest bit that's set in N
    N1=N - 1L ; n=1L
    while (n<N): n=n<<1L 
    n = n >> 1L

    # Rabin-Miller test
    for c in sieve[:7]:
        a=long(c) ; d=1L ; t=n
        while (t):  # Iterate over the bits in N1
            x=(d*d) % N
            if x==1L and d!=1L and d!=N1: return 0  # Square root of 1 found
            if N1 & t: d=(x*a) % N
            else: d=x
            t = t >> 1L
        if d!=1L: return 0
    return 1

# Small primes used for checking primality; these are all the primes
# less than 256.  This should be enough to eliminate 80% of all the odd
# numbers before needing to do a Rabin-Miller test at all.

sieve=[13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 
       71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 
       137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 
       197, 199, 211, 223, 227, 229, 233, 239, 241, 251 ]

# Key generation functions.

def generateQ(randfunc):
    "Generate a value of Q as suggested in Applied Cryptography"
    import sha
    S=randfunc(20)
    hash1=sha.new(S).digest()
    hash2=sha.new(long2str(str2long(S)+1)).digest()
    q = long(0)
    for i in range(0,20):
        c=ord(hash1[i])^ord(hash2[i])
        if i==0: c=c | 128
        if i==19: c= c | 1
        q=q*256+c
    while (not isPrime(q)):
        q=q+2
    if pow(2,159L)<q<pow(2,160L): return S, q
    raise error, 'Bad q value generated'
    
# Generate a DSA modulus with L bits
def generate(bits, randfunc, progress_func=None):
    import sha
    if bits<160: raise error, 'Key length <160 bits'
    obj=DSAobj()
    # Generate string S and prime q
    if progress_func: apply(progress_func, ('p,q\n',))
    while (1):
        S, obj.q = generateQ(randfunc)
        n=(bits-1)/160
        C, N, V = 0, 2, {}
	b=(obj.q >> 5) & 15
	powb=pow(long(2), b)
	powL1=pow(long(2), bits-1)
        while C<4096:
            for k in range(0, n+1):
		V[k]=str2long(sha.new(S+str(N)+str(k)).digest())
            W=V[n] % powb
            for k in range(n-1, -1, -1): W=(W<<160L)+V[k]
            X=W+powL1
            p=X-(X%(2*obj.q)-1)
            if powL1<=p and isPrime(p): break
            C, N = C+1, N+n+1
        if C<4096: break
	if progress_func: apply(progress_func, ('4096 multiples failed\n',) )
    obj.p = p
    power=(p-1)/obj.q
    if progress_func: apply(progress_func, ('h,g\n',))
    while (1):
        h=str2long(randfunc(bits)) % (p-1)
        g=pow(h, power, p)
        if 1<h<p-1 and g>1: break
    obj.g=g
    if progress_func: apply(progress_func, ('x,y\n',))
    while (1):
        x=str2long(randfunc(20))
        if 0<x<obj.q: break
    obj.x, obj.y=x, pow(g, x, p)
    return obj
    
def construct(tuple):
    obj=DSAobj()
    if len(tuple) not in [4,5]:
        raise error, 'argument for construct() wrong length' 
    for i in range(len(tuple)):
	field = obj.keydata[i]
	setattr(obj, field, tuple[i])
    return obj
    
class DSAobj:
    keydata=['y', 'g', 'p', 'q', 'x']

    def __init__(self):
	pass

    def sign(self, M, K):
	"""Return a tuple containing the signature for the message M.
        K is a random parameter required for generating a signature."""
	if (not self.hasprivate()):
	    raise error, 'Private key not available in this DSA object'
	if type(M)==types.StringType: M=str2long(M)
	if type(K)==types.StringType: K=str2long(K)
	if (K<2 or self.q<=K): raise error, 'K is not between 2 and q'
        r=pow(self.g, K, self.p) % self.q
        s=(inverse(K, self.q)*(M+self.x*r)) % self.q
        return (r,s)

    def verify(self, M, signature):
	"""Verify that the signature is valid for the message M;
	returns true if the signature checks out."""
	if type(M)==types.StringType: M=str2long(M)
        r, s = signature
	if r<=0 or r>=self.q or s<=0 or s>=self.q: return 0
        w=inverse(s, self.q)
        u1, u2 = (M*w) % self.q, (r*w) % self.q
        v1=pow(self.g, u1, self.p)
	v2=pow(self.y, u2, self.p)
	v=((v1*v2) % self.p)
	v=v % self.q
        if v==r: return 1
        return 0

    def hasprivate(self):
	"""Return a Boolean denoting whether the object contains
	private components."""
	if hasattr(self, 'x'): return 1
	else: return 0

    def publickey(self):
	"""Return a new key object containing only the public information."""
        return construct((self.y, self.g, self.p, self.q))

object=DSAobj


# Simple self-test

def testkey(key, randfunc):
    plaintext="Hello"

    print '  Signature test'
    K=getPrime(30, randfunc)
    signature=key.sign(plaintext, K)
    result=key.verify(plaintext, signature)
    if not result:
        print "***ERROR 1: Sig. verification failed when it should have succeeded"
    result=key.verify(plaintext[:-1], signature)
    if result:
        print "***ERROR 2: Sig. verification succeeded when it should have failed"
    # Change a single bit in the plaintext
    badtext=plaintext[:-3]+chr( 1 ^ ord(plaintext[-3]) )+plaintext[-3:]
    result=key.verify(badtext, signature)
    if result:
        print "***ERROR 3: Sig. verification succeeded when it should have failed"
    print '  Removing private key data'
    pubonly=key.publickey()
    result=pubonly.verify(plaintext, signature)
    if not result:
        print "***ERROR 4: Sig. verification failed when it should have succeeded"



if __name__ == '__main__':
    # This is a simple function that will generate numbers that look
    # vaguely random

    import time
    def randfunc(N, L = range(0,64)):
        s = ""
        L[0] = L[0] ^ (int(time.time()) % 256)
        N_orig = N
        while N>64:
            for i in range(64):
                s=s+chr(L[i])
                # Scramble the L array a bit
                L[i] = L[i] ^ L[ (i+1) % 64] ^ (L[ (i+8)%64] << 4)
                L[i] = L[i] ^ (int(time.time()) % 256) 
                L[i] = L[i] % 256
            N = N - 64
        for i in range(N): s=s+chr(L[i])
        assert len(s) == N_orig
        return s
        
    N=256				# Key size, measured in bits

    print ' Generating', N, 'bit key'
    import sys
    def write(s):
	import sys ; sys.stdout.write('  '+s)
    key=generate(N, randfunc, write)
    
    print ' Key data:'
    for field in key.keydata:
        print "  ", field, ':', hex(getattr(key,field))

    print " Testing newly generated key"
    testkey(key, randfunc)
    print " Testing pickled/unpickled key"
    import pickle
    s = pickle.dumps(key) ; key2 = pickle.loads(s)
    testkey(key2, randfunc)

    print " Testing cPickled key"
    import cPickle
    s = cPickle.dumps(key) ; key2 = cPickle.loads(s)
    testkey(key2, randfunc)
    print
    
