#!/usr/bin/env python

"""
    Helix: Fast Encryption and Authentication 
    by Niels Ferguson and Bruce Schneier

    Provides both ciphertext and an authentication tag with a 
    single pass of the data.  Typically, to get both two passes
    are required.

    Helix is also fast compared to other algorithms.

    See DDJ, Nov 2003, p28

    This code has been Psyco'd.

    Larry Bugbee
    October 2003
"""

try:
    import psyco
    _have_psyco = 1
    print 'psyco has been enabled'
except:
    _have_psyco = 0
    print 'psyco is not installed'

# struct provides conversion between strings and a list of words
import struct

def _strToWords( s ):
    # Interpret string as list of 32-bit integers, LSByte first.
    return list(struct.unpack( '<' + 'L' * (len(s)/4), s ))

def _wordsToStr( w ):
    # Convert list of 32-bit integers to a string, LSByte first.
    return apply( struct.pack, [ '<' + 'L' * len(w)] + w )

_mask32 = 0xffffffffL

def _rol32( V, n ):
    # Rotate a 32-bit value V left by n positions
    return _mask32 & ( (V<<n) | ((V & _mask32) >> (32-n) ) )

if _have_psyco:
    _rol32 = psyco.proxy(_rol32)

class Helix:

    def helixBlock( self, X0, P, X1 ):
        # Apply the Helix block function to self.Z using
        # X0, P, and X1 as the three input words.

        Z = self.Z

        Z[0] += Z[3]; Z[3] = _rol32( Z[3], 15 )
        Z[1] += Z[4]; Z[4] = _rol32( Z[4], 25 )
        Z[2] ^= Z[0]; Z[0] = _rol32( Z[0], 9 )
        Z[3] ^= Z[1]; Z[1] = _rol32( Z[1], 10 )
        Z[4] += Z[2]; Z[2] = _rol32( Z[2], 17 )

        Z[0] ^= (Z[3] + X0); Z[3] = _rol32( Z[3], 30 )
        Z[1] ^= Z[4]; Z[4] = _rol32( Z[4], 13 )
        Z[2] += Z[0]; Z[0] = _rol32( Z[0], 20 )
        Z[3] += Z[1]; Z[1] = _rol32( Z[1], 11 )
        Z[4] ^= Z[2]; Z[2] = _rol32( Z[2], 5 )

        Z[0] += (Z[3] ^ P); Z[3] = _rol32( Z[3], 15 )
        Z[1] += Z[4]; Z[4] = _rol32( Z[4], 25 )
        Z[2] ^= Z[0]; Z[0] = _rol32( Z[0], 9 )
        Z[3] ^= Z[1]; Z[1] = _rol32( Z[1], 10 )
        Z[4] += Z[2]; Z[2] = _rol32( Z[2], 17 )

        Z[0] ^= (Z[3] + X1); Z[3] = _rol32( Z[3], 30 )
        Z[1] ^= Z[4]; Z[4] = _rol32( Z[4], 13 )
        Z[2] += Z[0]; Z[0] = _rol32( Z[0], 20 )
        Z[3] += Z[1]; Z[1] = _rol32( Z[1], 11 )
        Z[4] ^= Z[2]; Z[2] = _rol32( Z[2], 5 )

    if _have_psyco:
        helixBlock = psyco.proxy(helixBlock)

    def __init__( self, key ):
        # Initialise new Helix object with a key.
        # The key is a string with length <= 32.   (32x8 = 256 bits)
        self.lK = len( key )
        assert 0 <= self.lK <= 32
        # Convert to list of 8 words
        self.K = _strToWords( key + '\000'*(32-self.lK) )
        # Perform the key mixing
        for i in range( 8 ):
            self.Z = self.K[:4] + [self.lK + 64]
            self.helixBlock( 0, 0, 0 )
            self.K = [self.K[4]^self.Z[0],
                self.K[5]^self.Z[1],
                self.K[6]^self.Z[2],
                self.K[7]^self.Z[3]] + self.K[:4]

    def doBlk( self, P ):
        # Perform single Helix block operation with plaintext word P
        # computing X_i,0 and X_i,1 appropriately.
        # Get i mod 8
        i = self.i8 % 8
        # Compute X1 from self.X1 and (i+8)
        X1 = self.X1[ i ]
        if i%4 == 3:
            X1 += (self.i8)>>31
        X1 = (X1 + self.i8) & _mask32
        # The actual block function
        self.helixBlock( self.K[i], P, X1 )
        # Update the block number
        self.i8 += 1

    if _have_psyco:
        doBlk = psyco.proxy(doBlk)

    def start( self, N ):
        # Initialise a Helix encryption or decryption with nonce N.
        # The nonce is a string of length 16.   (16x8 = 128 bits)
        assert len(N) == 16
        Nl = _strToWords( N )
        # Extend the nonce to 8 words
        for i in range( 4 ):
            Nl.append( (i - Nl[i]) & _mask32 )
        K = self.K
        # Make array of X1 values to generate X_{i,1} efficiently.
        self.X1 = []
        for i in range( 8 ):
            x = ((i%4)==1)*4*self.lK
            self.X1.append( (K[(i+4)%8] + Nl[i] + x) & _mask32 )
        # Initialise state and run first 8 blocks for nonce mixing.
        self.Z = [K[3]^Nl[0], K[4]^Nl[1], K[5]^Nl[2], K[6]^Nl[3], K[7] ]
        self.i8 = 0
        for i in range( 8 ):
            self.doBlk( 0 )

    def finish( self, lnm4 ):
        # Finish up Helix processing and return the Tag.
        # Tweak the internal state as specified
        self.Z[0] ^= 0x912d94f1L
        # Apply 8 block functions with len(P) mod 4 as plaintext
        for i in range( 8 ):
            self.doBlk( lnm4 )
        # And generate the tag
        tag = []
        for i in range( 4 ):
            tag.append( self.Z[0] )
            self.doBlk( lnm4 )
        # Cleanup.
        self.Z = None
        self.X1 = None
        self.i8 = None
        # return the tag as a string of bytes
        return _wordsToStr( tag )

    def encrypt( self, N, P ):
        """
        Encrypt the plaintext P with the nonce N.
        Returns a tuple (C,T) containing the ciphertext C and the
        authentication tag T.
        """
        # Initialise a Helix encryption
        self.start( N )
        ln = len(P)
        # Pad plaintext with zeroes to make a whole number of words
        P = P + '\000' * ( (4-(ln%4)) % 4 )
        # ... and convert P to a list of 32-bit words.
        Pl = _strToWords( P )
        # Cl will contain the ciphertext, encoded as a list of words.
        Cl = []
        # Main encryption loop
        for p in Pl:
            Cl.append( p ^ self.Z[0] )
            self.doBlk( p )
        # Convert ciphertext to string, and trim to proper length
        C = _wordsToStr( Cl )[:ln]
        # Compute the authentication tag (requires len(P) mod 4).
        T = self.finish( ln%4 )
        return (C,T)

    def decrypt( self, N, C, T ):
        """
        Decrypt the plaintext C with the nonce N, and check that the
        authentication tag T is correct.
        Returns None if the authentication failed;
        returns the plaintext P if the authentication was correct.
        """
        # Initialise Helix decryption with the nonce.
        self.start( N )
        # Compute the ciphertext words and the mask words. The last plaintext word
        # needs to be masked to ensure that unused plaintext bytes are set to zero
        # for the tag computation. We mask every word, which is slightly easier
        # to implement though it is less efficient.
        ln = len( C )
        # Pad C and compute mask string (in bytes)
        pad = '\000' * ( (4-(ln%4)) % 4 )
        C = C + pad
        M = '\xff' * ln + pad
        # Convert to words
        Cl = _strToWords( C )
        Ml = _strToWords( M )
        # Pl will contain the plaintext as a list of words.
        Pl = []
        # Main decryption loop
        for i in range( len( Cl ) ):
            # Decrypt a word, and apply the mask
            p = (self.Z[0] ^ Cl[i]) & Ml[i]
            Pl.append( p )
            self.doBlk( p )
        # Convert plaintext to bytes, and trim to correct length.
        P = _wordsToStr( Pl )[:ln]
        # Compute the value we expect the authentication tag to have.
        V = self.finish( ln%4 )
        # Check for the correct tag value
        if V == T:
            return P
        else:
            # The tag was incorrect.
            return None



if __name__ == '__main__':
    print '-'*40
    h = Helix('myKey')          # must be a string of <= 32 char
    N = '1234567890123456'      # must be a string of exactly 16 char (128 bits)
    C, T = h.encrypt(N, 'Kilroy was here! ...and now is the time for everything.')
    P = h.decrypt(N, C, T)
    print P
    print '-'*40