pureSalsa20.py

#!/usr/bin/env python
# coding: utf-8

"""
    Copyright by https://github.com/zhansliu/writemdict

    pureSalsa20.py -- a pure Python implementation of the Salsa20 cipher, ported to Python 3

    v4.0: Added Python 3 support, dropped support for Python <= 2.5.
    
    // zhansliu

    Original comments below.

    ====================================================================
    There are comments here by two authors about three pieces of software:
        comments by Larry Bugbee about
            Salsa20, the stream cipher by Daniel J. Bernstein 
                 (including comments about the speed of the C version) and
            pySalsa20, Bugbee's own Python wrapper for salsa20.c
                 (including some references), and
        comments by Steve Witham about
            pureSalsa20, Witham's pure Python 2.5 implementation of Salsa20,
                which follows pySalsa20's API, and is in this file.

    Salsa20: a Fast Streaming Cipher (comments by Larry Bugbee)
    -----------------------------------------------------------

    Salsa20 is a fast stream cipher written by Daniel Bernstein 
    that basically uses a hash function and XOR making for fast 
    encryption.  (Decryption uses the same function.)  Salsa20 
    is simple and quick.  
    
    Some Salsa20 parameter values...
        design strength    128 bits
        key length         128 or 256 bits, exactly
        IV, aka nonce      64 bits, always
        chunk size         must be in multiples of 64 bytes
    
    Salsa20 has two reduced versions, 8 and 12 rounds each.
    
    One benchmark (10 MB):
        1.5GHz PPC G4     102/97/89 MB/sec for 8/12/20 rounds
        AMD Athlon 2500+   77/67/53 MB/sec for 8/12/20 rounds
          (no I/O and before Python GC kicks in)
    
    Salsa20 is a Phase 3 finalist in the EU eSTREAM competition 
    and appears to be one of the fastest ciphers.  It is well 
    documented so I will not attempt any injustice here.  Please 
    see "References" below.
    
    ...and Salsa20 is "free for any use".  
    
    
    pySalsa20: a Python wrapper for Salsa20 (Comments by Larry Bugbee)
    ------------------------------------------------------------------

    pySalsa20.py is a simple ctypes Python wrapper.  Salsa20 is 
    as it's name implies, 20 rounds, but there are two reduced 
    versions, 8 and 12 rounds each.  Because the APIs are 
    identical, pySalsa20 is capable of wrapping all three 
    versions (number of rounds hardcoded), including a special 
    version that allows you to set the number of rounds with a 
    set_rounds() function.  Compile the version of your choice 
    as a shared library (not as a Python extension), name and 
    install it as libsalsa20.so.
    
    Sample usage:
        from pySalsa20 import Salsa20
        s20 = Salsa20(key, IV)
        dataout = s20.encryptBytes(datain)   # same for decrypt
    
    This is EXPERIMENTAL software and intended for educational 
    purposes only.  To make experimentation less cumbersome, 
    pySalsa20 is also free for any use.      
    
    THIS PROGRAM IS PROVIDED WITHOUT WARRANTY OR GUARANTEE OF
    ANY KIND.  USE AT YOUR OWN RISK.  
    
    Enjoy,
      
    Larry Bugbee
    bugbee@seanet.com
    April 2007

    
    References:
    -----------
      http://en.wikipedia.org/wiki/Salsa20
      http://en.wikipedia.org/wiki/Daniel_Bernstein
      http://cr.yp.to/djb.html
      http://www.ecrypt.eu.org/stream/salsa20p3.html
      http://www.ecrypt.eu.org/stream/p3ciphers/salsa20/salsa20_p3source.zip

     
    Prerequisites for pySalsa20:
    ----------------------------
      - Python 2.5 (haven't tested in 2.4)


    pureSalsa20: Salsa20 in pure Python 2.5 (comments by Steve Witham)
    ------------------------------------------------------------------

    pureSalsa20 is the stand-alone Python code in this file.
    It implements the underlying Salsa20 core algorithm
    and emulates pySalsa20's Salsa20 class API (minus a bug(*)).

    pureSalsa20 is MUCH slower than libsalsa20.so wrapped with pySalsa20--
    about 1/1000 the speed for Salsa20/20 and 1/500 the speed for Salsa20/8,
    when encrypting 64k-byte blocks on my computer.

    pureSalsa20 is for cases where portability is much more important than
    speed.  I wrote it for use in a "structured" random number generator.

    There are comments about the reasons for this slowness in
          http://www.tiac.net/~sw/2010/02/PureSalsa20

    Sample usage:
        from pureSalsa20 import Salsa20
        s20 = Salsa20(key, IV)
        dataout = s20.encryptBytes(datain)   # same for decrypt

    I took the test code from pySalsa20, added a bunch of tests including
    rough speed tests, and moved them into the file testSalsa20.py.  
    To test both pySalsa20 and pureSalsa20, type
        python testSalsa20.py

    (*)The bug (?) in pySalsa20 is this.  The rounds variable is global to the
    libsalsa20.so library and not switched when switching between instances
    of the Salsa20 class.
        s1 = Salsa20( key, IV, 20 )
        s2 = Salsa20( key, IV, 8 )
    In this example,
        with pySalsa20, both s1 and s2 will do 8 rounds of encryption.
        with pureSalsa20, s1 will do 20 rounds and s2 will do 8 rounds.
    Perhaps giving each instance its own nRounds variable, which
    is passed to the salsa20wordtobyte() function, is insecure.  I'm not a 
    cryptographer.

    pureSalsa20.py and testSalsa20.py are EXPERIMENTAL software and 
    intended for educational purposes only.  To make experimentation less 
    cumbersome, pureSalsa20.py and testSalsa20.py are free for any use.

    Revisions:
    ----------
      p3.2   Fixed bug that initialized the output buffer with plaintext!
             Saner ramping of nreps in speed test.
             Minor changes and print statements.
      p3.1   Took timing variability out of add32() and rot32().
             Made the internals more like pySalsa20/libsalsa .
             Put the semicolons back in the main loop!
             In encryptBytes(), modify a byte array instead of appending.
             Fixed speed calculation bug.
             Used subclasses instead of patches in testSalsa20.py .
             Added 64k-byte messages to speed test to be fair to pySalsa20.
      p3     First version, intended to parallel pySalsa20 version 3.

    More references:
    ----------------
      http://www.seanet.com/~bugbee/crypto/salsa20/          [pySalsa20]
      http://cr.yp.to/snuffle.html        [The original name of Salsa20]
      http://cr.yp.to/snuffle/salsafamily-20071225.pdf [ Salsa20 design]
      http://www.tiac.net/~sw/2010/02/PureSalsa20
    
    THIS PROGRAM IS PROVIDED WITHOUT WARRANTY OR GUARANTEE OF
    ANY KIND.  USE AT YOUR OWN RISK.  

    Cheers,

    Steve Witham sw at remove-this tiac dot net
    February, 2010
"""
import sys
assert(sys.version_info >= (2, 6))

if sys.version_info >= (3,):
	integer_types = (int,)
	python3 = True
else:
	integer_types = (int, long)
	python3 = False

from struct import Struct
little_u64 = Struct( "<Q" )      #    little-endian 64-bit unsigned.
                                 #    Unpacks to a tuple of one element!

little16_i32 = Struct( "<16i" )  # 16 little-endian 32-bit signed ints.
little4_i32 = Struct( "<4i" )    #  4 little-endian 32-bit signed ints.
little2_i32 = Struct( "<2i" )    #  2 little-endian 32-bit signed ints.

_version = 'p4.0'

#----------- Salsa20 class which emulates pySalsa20.Salsa20 ---------------

class Salsa20(object):
    def __init__(self, key=None, IV=None, rounds=20 ):
        self._lastChunk64 = True
        self._IVbitlen = 64             # must be 64 bits
        self.ctx = [ 0 ] * 16
        if key:
            self.setKey(key)
        if IV:
            self.setIV(IV)

        self.setRounds(rounds)


    def setKey(self, key):
        assert type(key) == bytes
        ctx = self.ctx
        if len( key ) == 32:  # recommended
            constants = b"expand 32-byte k"
            ctx[ 1],ctx[ 2],ctx[ 3],ctx[ 4] = little4_i32.unpack(key[0:16])
            ctx[11],ctx[12],ctx[13],ctx[14] = little4_i32.unpack(key[16:32])
        elif len( key ) == 16:
            constants = b"expand 16-byte k"
            ctx[ 1],ctx[ 2],ctx[ 3],ctx[ 4] = little4_i32.unpack(key[0:16])
            ctx[11],ctx[12],ctx[13],ctx[14] = little4_i32.unpack(key[0:16])
        else:
            raise Exception( "key length isn't 32 or 16 bytes." )
        ctx[0],ctx[5],ctx[10],ctx[15] = little4_i32.unpack( constants )

        
    def setIV(self, IV):
        assert type(IV) == bytes
        assert len(IV)*8 == 64, 'nonce (IV) not 64 bits'
        self.IV = IV
        ctx=self.ctx
        ctx[ 6],ctx[ 7] = little2_i32.unpack( IV )
        ctx[ 8],ctx[ 9] = 0, 0  # Reset the block counter.

    setNonce = setIV            # support an alternate name


    def setCounter( self, counter ):
        assert( type(counter) in integer_types )
        assert( 0 <= counter < 1<<64 ), "counter < 0 or >= 2**64"
        ctx = self.ctx
        ctx[ 8],ctx[ 9] = little2_i32.unpack( little_u64.pack( counter ) )

    def getCounter( self ):
        return little_u64.unpack( little2_i32.pack( *self.ctx[ 8:10 ] ) ) [0]


    def setRounds(self, rounds, testing=False ):
        assert testing or rounds in [8, 12, 20], 'rounds must be 8, 12, 20'
        self.rounds = rounds


    def encryptBytes(self, data):
        assert type(data) == bytes, 'data must be byte string'
        assert self._lastChunk64, 'previous chunk not multiple of 64 bytes'
        lendata = len(data)
        munged = bytearray(lendata)
        for i in range( 0, lendata, 64 ):
            h = salsa20_wordtobyte( self.ctx, self.rounds, checkRounds=False )
            self.setCounter( ( self.getCounter() + 1 ) % 2**64 )
            # Stopping at 2^70 bytes per nonce is user's responsibility.
            for j in range( min( 64, lendata - i ) ):
                if python3:
                    munged[ i+j ] = data[ i+j ] ^ h[j]
                else:
                    munged[ i+j ] = ord(data[ i+j ]) ^ ord(h[j])

        self._lastChunk64 = not lendata % 64
        return bytes(munged)
    
    decryptBytes = encryptBytes # encrypt and decrypt use same function

#--------------------------------------------------------------------------

def salsa20_wordtobyte( input, nRounds=20, checkRounds=True ):
    """ Do nRounds Salsa20 rounds on a copy of 
            input: list or tuple of 16 ints treated as little-endian unsigneds.
        Returns a 64-byte string.
        """

    assert( type(input) in ( list, tuple )  and  len(input) == 16 )
    assert( not(checkRounds) or ( nRounds in [ 8, 12, 20 ] ) )

    x = list( input )

    def XOR( a, b ):  return a ^ b
    ROTATE = rot32
    PLUS   = add32

    for i in range( nRounds // 2 ):
        # These ...XOR...ROTATE...PLUS... lines are from ecrypt-linux.c
        # unchanged except for indents and the blank line between rounds:
        x[ 4] = XOR(x[ 4],ROTATE(PLUS(x[ 0],x[12]), 7));
        x[ 8] = XOR(x[ 8],ROTATE(PLUS(x[ 4],x[ 0]), 9));
        x[12] = XOR(x[12],ROTATE(PLUS(x[ 8],x[ 4]),13));
        x[ 0] = XOR(x[ 0],ROTATE(PLUS(x[12],x[ 8]),18));
        x[ 9] = XOR(x[ 9],ROTATE(PLUS(x[ 5],x[ 1]), 7));
        x[13] = XOR(x[13],ROTATE(PLUS(x[ 9],x[ 5]), 9));
        x[ 1] = XOR(x[ 1],ROTATE(PLUS(x[13],x[ 9]),13));
        x[ 5] = XOR(x[ 5],ROTATE(PLUS(x[ 1],x[13]),18));
        x[14] = XOR(x[14],ROTATE(PLUS(x[10],x[ 6]), 7));
        x[ 2] = XOR(x[ 2],ROTATE(PLUS(x[14],x[10]), 9));
        x[ 6] = XOR(x[ 6],ROTATE(PLUS(x[ 2],x[14]),13));
        x[10] = XOR(x[10],ROTATE(PLUS(x[ 6],x[ 2]),18));
        x[ 3] = XOR(x[ 3],ROTATE(PLUS(x[15],x[11]), 7));
        x[ 7] = XOR(x[ 7],ROTATE(PLUS(x[ 3],x[15]), 9));
        x[11] = XOR(x[11],ROTATE(PLUS(x[ 7],x[ 3]),13));
        x[15] = XOR(x[15],ROTATE(PLUS(x[11],x[ 7]),18));

        x[ 1] = XOR(x[ 1],ROTATE(PLUS(x[ 0],x[ 3]), 7));
        x[ 2] = XOR(x[ 2],ROTATE(PLUS(x[ 1],x[ 0]), 9));
        x[ 3] = XOR(x[ 3],ROTATE(PLUS(x[ 2],x[ 1]),13));
        x[ 0] = XOR(x[ 0],ROTATE(PLUS(x[ 3],x[ 2]),18));
        x[ 6] = XOR(x[ 6],ROTATE(PLUS(x[ 5],x[ 4]), 7));
        x[ 7] = XOR(x[ 7],ROTATE(PLUS(x[ 6],x[ 5]), 9));
        x[ 4] = XOR(x[ 4],ROTATE(PLUS(x[ 7],x[ 6]),13));
        x[ 5] = XOR(x[ 5],ROTATE(PLUS(x[ 4],x[ 7]),18));
        x[11] = XOR(x[11],ROTATE(PLUS(x[10],x[ 9]), 7));
        x[ 8] = XOR(x[ 8],ROTATE(PLUS(x[11],x[10]), 9));
        x[ 9] = XOR(x[ 9],ROTATE(PLUS(x[ 8],x[11]),13));
        x[10] = XOR(x[10],ROTATE(PLUS(x[ 9],x[ 8]),18));
        x[12] = XOR(x[12],ROTATE(PLUS(x[15],x[14]), 7));
        x[13] = XOR(x[13],ROTATE(PLUS(x[12],x[15]), 9));
        x[14] = XOR(x[14],ROTATE(PLUS(x[13],x[12]),13));
        x[15] = XOR(x[15],ROTATE(PLUS(x[14],x[13]),18));

    for i in range( len( input ) ):
        x[i] = PLUS( x[i], input[i] )
    return little16_i32.pack( *x )

#--------------------------- 32-bit ops -------------------------------

def trunc32( w ):
    """ Return the bottom 32 bits of w as a Python int.
        This creates longs temporarily, but returns an int. """
    w = int( ( w & 0x7fffFFFF ) | -( w & 0x80000000 ) )
    assert type(w) == int
    return w


def add32( a, b ):
    """ Add two 32-bit words discarding carry above 32nd bit,
        and without creating a Python long.
        Timing shouldn't vary.
    """
    lo = ( a & 0xFFFF ) + ( b & 0xFFFF )
    hi = ( a >> 16 ) + ( b >> 16 ) + ( lo >> 16 )
    return ( -(hi & 0x8000) | ( hi & 0x7FFF ) ) << 16 | ( lo & 0xFFFF )


def rot32( w, nLeft ):
    """ Rotate 32-bit word left by nLeft or right by -nLeft
        without creating a Python long.
        Timing depends on nLeft but not on w.
    """
    nLeft &= 31  # which makes nLeft >= 0
    if nLeft == 0:
        return w

    # Note: now 1 <= nLeft <= 31.
    #     RRRsLLLLLL   There are nLeft RRR's, (31-nLeft) LLLLLL's,
    # =>  sLLLLLLRRR   and one s which becomes the sign bit.
    RRR = ( ( ( w >> 1 ) & 0x7fffFFFF ) >> ( 31 - nLeft ) )
    sLLLLLL = -( (1<<(31-nLeft)) & w ) | (0x7fffFFFF>>nLeft) & w
    return RRR | ( sLLLLLL << nLeft )


# --------------------------------- end -----------------------------------