chachapoly.py

import array
import hashlib
import secrets
import struct


__all__ = ['ChaCha']


class ChaCha20Poly1305Aead():
    """ https://datatracker.ietf.org/doc/html/rfc8439 """
    
    def __init__(self, key=None):
        key = key or secrets.token_bytes(32)
        if not isinstance(key, bytes):
            raise TypeError('Key must be 32 {bytes!r}. Got: {key!r}')
        if len(key) != 32:
            raise ValueError('Key must be 32 {bytes!r}. Got {key!r}')
        
        self._key = key
    
    def encryptor(self, iv=None, **kwargs):
        iv = iv or secrets.token_bytes(12)
        if len(iv) != 12:
            raise ValueError('iv must be 12 bytes')

        return _ChaChaPolyAeadEncryptor(self._key, iv, **kwargs)
        
        
    def decryptor(self, iv=None, **kwargs):
        iv = iv or secrets.token_bytes(12)
        if len(iv) != 12:
            raise ValueError('iv must be 12 bytes')

        return _ChaChaPolyAeadDecryptor(self._key, iv, **kwargs)
        
    
    def cipher_pair(self, iv=None, **kwargs):
        iv = iv or secrets.token_bytes(12)
        if len(iv) != 12:
            raise ValueError('iv must be 12 bytes')
            
        e = _ChaChaPolyAeadEncryptor(self._key, iv, **kwargs)
        d = _ChaChaPolyAeadDecryptor(self._key, iv, **kwargs)
        e.set_ratchet_friend(d)
        
        return e, d


class _Poly1305Mac():
    def __init__(self, key:bytes):
        if len(key) != 32:
            raise ValueError('Key must be 32 bytes')
        
        self._r = int.from_bytes(key[:16], 'little')
        self._s = int.from_bytes(key[16:], 'little')
        # magic number is bit clamp per RFC
        self._r &= 0x0ffffffc0ffffffc0ffffffc0fffffff
        self._a = 0
        self._remainder = b''
        

    def _update(self, msg):
        a = int.from_bytes(msg, 'little')
        a |=  2**(len(msg) * 8)
        # magic number is prime P = 2^130-5 per RFC
        return ((self._a + a) * self._r) % 0x3fffffffffffffffffffffffffffffffb
        
    
    def update(self, message:bytes):
        message = self._remainder + message
        if len(message) < 16:
            self._remainder = message
        elif len(message) == 16:
            self._remainder = b''
            self._a = self._update(message)
        else:
            for i in range(0, len(message)-16, 16):
                self._a = self._update(message[i:i+16])
            self._remainder = message[i+16:]
        
    
    def digest(self):
        if self._remainder:
            tag = self._update(self._remainder)
        else:
            tag = self._a
        tag = (tag + self._s) & 0xffffffffffffffffffffffffffffffff
        return tag.to_bytes(16, 'little')
    
    
    def hexdigest(self):
        return self.digest().hex()
    


class _ChaCha20Cipher():
    def __init__(self, key, iv, counter=0):
        self._key = key
        self._iv = iv
        s = b'expand 32-byte k' + key + struct.pack('<I', counter) + iv
        if len(s) != 64:
            raise ValueError(
                f'Invalid buffer size. Expecting: {64-len(s)}'
            )
            
        a = array.array('L')
        a.frombytes(s)  # this actually converts to little-endian
        if a[0] == 0x65787061:
            a.byteswap()
        if a[0] != 0x61707865:
            raise ValueError('ChaCha cipher could not be initialized properly')
        self._state = a  # only update the counter on this (index 12)
        self._working_state = a.__copy__()  # this gets chacha'd
        self._counter = counter
        
    
    def _quarter_round(self, a, b, c, d):
        s = self._working_state
        s[a]=(s[a]+s[b])&0xffffffff; s[d]^=s[a]; s[d]=((s[d]<<16)|(s[d]>>16))&0xffffffff
        s[c]=(s[c]+s[d])&0xffffffff; s[b]^=s[c]; s[b]=((s[b]<<12)|(s[b]>>20))&0xffffffff
        s[a]=(s[a]+s[b])&0xffffffff; s[d]^=s[a]; s[d]=((s[d]<<8 )|(s[d]>>24))&0xffffffff
        s[c]=(s[c]+s[d])&0xffffffff; s[b]^=s[c]; s[b]=((s[b]<<7 )|(s[b]>>25))&0xffffffff
    
    
    def _half_round(self):
        self._quarter_round(0, 4, 8 , 12)
        self._quarter_round(1, 5, 9 , 13)
        self._quarter_round(2, 6, 10, 14)
        self._quarter_round(3, 7, 11, 15)
        self._quarter_round(0, 5, 10, 15)
        self._quarter_round(1, 6, 11, 12)
        self._quarter_round(2, 7, 8 , 13)
        self._quarter_round(3, 4, 9 , 14)
  
  
    def _full_round(self):
        for _ in range(10):
            self._half_round()
        
        
    def _key_stream(self):
        for c in range(0x100000000):
            c = (c + self._counter) & 0xffffffff  # wrap 32-bit int 
            self._state[12] = c
            self._working_state[:] = self._state
            self._full_round()
            
            for i in range(16):
                self._working_state[i] = (
                    self._state[i] + self._working_state[i]
                ) & 0xffffffff
            
            for b in self._working_state.tobytes():
                yield b
            
        
        raise RuntimeError('Keyspace exhausted')
    
    
    @property
    def _stream(self):
        """Stream of 8-bit integers"""
        return self._key_stream()
    
    
    @property
    def blockid(self):
        return self._counter
    


class _PFSMixin():
    """Mixin for Perfect Forward Secrecy. Seems to work okay..."""
    
    def set_ratchet_friend(self, friend):
        """Interface to simplify ratchet synchronization"""
        if isinstance(friend, _ChaChaPolyAeadCryptor):
            self._ratchet_friend = friend
            friend._ratchet_friend = self
        
    
    def ratchet(self, key):
        """Used with a DH protocol to provide forward secrecy.
        Input: Random temporal key 
        
        No output, automatically configures internal object's cipher state
        with new 256-bit key and 96-bit nonce.
        """
        kdf_key = getattr(self, '_kdf_key', b'')
        data = hashlib.scrypt(
            key + kdf_key,
            salt=self._state.tobytes(),
            n=2**18,
            r=8,
            p=1,
            maxmem=0x7fffffff,
            dklen=32 + 32 + 12
        )
        self._key = data[:32]
        key = array.array('L')
        key.frombytes(self._key)
        self._state[4:12] = key
        
        self._iv = data[32:44]
        iv = array.array('L')
        iv.frombytes(self._iv)
        self._state[13:16] = iv
        
        self._kdf_key = data[44:76]
        
        friend = getattr(self, '_ratchet_friend', None)
        if friend is not None:
            friend._key = self._key
            friend._state[4:12] = key
            friend._iv = self._iv
            friend._state[13:16] = iv
            friend._kdf_key = self._kdf_key
        
    
    

class _ChaChaPolyAeadCryptor(_ChaCha20Cipher, _PFSMixin):
    """ https://datatracker.ietf.org/doc/html/rfc8439#section-2.6 """
    def __init__(self, key, iv, auth=None, counter=0, mac_key=None):
        super().__init__(key, iv, counter)
        
        # Initialize the Poly1305 Mac
        if mac_key:
            self._mac = _Poly1305Mac(mac_key)
        else:
            a = self._stream
            mac_key = bytes(next(a) for _ in range(32))
            self._mac = _Poly1305Mac(mac_key)

        self._counter = 1
        
        if auth is not None:
            if not isinstance(auth, bytes):
                raise TypeError(f'Auth data must be bytes, not {auth!r}')
        
        self._auth = auth or b''
        self._bytes_processed = 0
    
    def finalize(self, data=None):
        """Compute and return any remaining encryption data with HMAC 
        
            ciphertext format: Ciphertext + tag
        
        """
        raise NotImplementedError
    
    def update(self, data):
        """Encryption per https://datatracker.ietf.org/doc/html/rfc8439#section-2.8 """
        raise NotImplementedError
    
    def encrypt(self, data):
        """One-shot encryption. Can only be called once"""
        raise NotImplementedError
    
    def decrypt(self, data):
        """One-shot decryption. Can only be called once"""
        raise NotImplementedError
    
    @property
    def key(self):
        return self._key
    
    @property
    def nonce(self):
        return self._iv
    


class _ChaChaPolyAeadEncryptor(_ChaChaPolyAeadCryptor):
    def finalize(self, data=None):
        if data:
            ct = self.update(data)
        else:
            ct = b''
        if (pad:=self._bytes_processed % 16) > 0:
            self._mac.update(b'\0' * (16 - pad))
        self._mac.update(struct.pack('<Q', len(self._auth)))
        self._mac.update(struct.pack('<Q', self._bytes_processed))
        self._bytes_processed = -1
        return ct + self._mac.digest()
        
    
    def update(self, data):
        if self._bytes_processed < 0:
            raise RuntimeError('Cannot update after .finalize()')
        ct = b''
        if data:
            if self._bytes_processed == 0:
                m = self._auth
                if len(m) % 16 > 0:
                    m += b'\0' * (16 - (len(m)%16))
                self._mac.update(m)
            ct = bytes(a^b for a,b in zip(self._stream, data))
            self._bytes_processed += len(ct)
            self._mac.update(ct)
        return ct
        
        
    def encrypt(self, data):
        if self._bytes_processed > 0:
            raise RuntimeError('Attempted nonce reuse (make a new .encryptor() instead)')
        return self.update(data) + self.finalize()
        
        
        
class _ChaChaPolyAeadDecryptor(_ChaChaPolyAeadCryptor):
    def finalize(self, data):
        
        if len(data) < 16:
            raise ValueError('Invalid MAC length')
        elif len(data) > 16:
            data, mac = data[:-16], data[-16:]
        else:
            data, mac = b'', data
        
        if data:
            pt = self.update(data)
        else:
            pt = b''
        if (pad:=self._bytes_processed % 16) > 0:
            self._mac.update(b'\0' * (16 - pad))
        self._mac.update(struct.pack('<Q', len(self._auth)))
        self._mac.update(struct.pack('<Q', self._bytes_processed))
        if mac == self._mac.digest():
            self._bytes_processed = -1
            return pt 
        raise ValueError('Invalid MAC')

    
    def update(self, data):
        if self._bytes_processed < 0:
            raise RuntimeError('Cannot update after .finalize()')
        pt = b''
        if data:
            if self._bytes_processed == 0:
                m = self._auth
                if len(m) % 16 > 0:
                    m += b'\0' * (16 - (len(m)%16))
                self._mac.update(m)
            self._mac.update(data)
            pt = bytes(a^b for a,b in zip(self._stream, data))
            self._bytes_processed += len(pt)
        return pt

    
    def decrypt(self, data):
        # assume encoded mac
        data, mac = data[:-16], data[-16:]
        if self._bytes_processed > 0:
            raise RuntimeError('Attempted mac reuse (make a new .decryptor() instead)')
        return self.update(data) + self.finalize(mac)
        


class ChaCha(ChaCha20Poly1305Aead):
    """Alias to simplify stuff"""
        
    
    


# magical tests from the RFC

if __name__ == '__main__':
    import itertools
    import unittest
    from unittest.mock import Mock
    
    class TestChaCha20Cipher(unittest.TestCase):
        def setUp(self):
            self.key = bytes.fromhex(
                '000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f'
            )
            self.iv = bytes.fromhex('000000090000004a00000000')
            self.cipher = _ChaCha20Cipher(self.key, self.iv, counter=1)


        def test_quarter_round(self):
            c = Mock()
            c._working_state = array.array('L')
            
            # bytes are in little endian, but bytes.fromhex interprets it as big
            c._working_state.frombytes(bytes.fromhex('11111111010203049b8d6f4301234567'))
            
            # swap the bytes... and now it's in little endian
            c._working_state.byteswap()
            _ChaCha20Cipher._quarter_round(c, 0, 1, 2, 3)
            
            # Since ea2a...c4bb is "big endian" again. swap again
            c._working_state.byteswap()
            
            self.assertEqual(
                c._working_state.tobytes(),
                bytes.fromhex('ea2a92f4cb1cf8ce4581472e5881c4bb')
            )
          
        
        
        def test_state_initialization(self):
            self.cipher._working_state.byteswap()  # compare in big endian
            self.assertEqual(
                self.cipher._working_state.tobytes(),
                bytes.fromhex('''
                    61707865  3320646e  79622d32  6b206574
                    03020100  07060504  0b0a0908  0f0e0d0c
                    13121110  17161514  1b1a1918  1f1e1d1c
                    00000001  09000000  4a000000  00000000
                ''')
            )
        
        
        def test_full_round_pre_mix(self):
            self.cipher._full_round()
            
            self.cipher._working_state.byteswap()  # compare in big endian
            self.assertEqual(
                self.cipher._working_state.tobytes(),
                bytes.fromhex('''
                    837778ab  e238d763  a67ae21e  5950bb2f
                    c4f2d0c7  fc62bb2f  8fa018fc  3f5ec7b7
                    335271c2  f29489f3  eabda8fc  82e46ebd
                    d19c12b4  b04e16de  9e83d0cb  4e3c50a2
                ''')
            )
        
        
        def test_full_round_post_mix(self):
            initial_state = self.cipher._working_state.__copy__()
            self.cipher._full_round()
            for i in range(16):
                self.cipher._working_state[i] = (
                    self.cipher._working_state[i] + initial_state[i]
                ) & 0xffffffff
            
            self.cipher._working_state.byteswap()  # compare in big endian
            self.assertEqual(
                self.cipher._working_state.tobytes(),
                bytes.fromhex('''
                    e4e7f110  15593bd1  1fdd0f50  c47120a3
                    c7f4d1c7  0368c033  9aaa2204  4e6cd4c3
                    466482d2  09aa9f07  05d7c214  a2028bd9
                    d19c12b5  b94e16de  e883d0cb  4e3c50a2
                ''')
            )
            
            
        def test_serialization(self):
            key_stream = self.cipher._stream
            first_64_bytes = bytes(next(key_stream) for _ in range(64))
            self.assertEqual(
                first_64_bytes,
                bytes.fromhex('''
                    10 f1 e7 e4 d1 3b 59 15 50 0f dd 1f a3 20 71 c4
                    c7 d1 f4 c7 33 c0 68 03 04 22 aa 9a c3 d4 6c 4e
                    d2 82 64 46 07 9f aa 09 14 c2 d7 05 d9 8b 02 a2
                    b5 12 9c d1 de 16 4e b9 cb d0 83 e8 a2 50 3c 4e
                ''')
            )
        


    class TestPoly1305(unittest.TestCase):
        def setUp(self):
            self.key = bytes.fromhex(
                '85d6be7857556d337f4452fe42d506a8'
                '0103808afb0db2fd4abff6af4149f51b'
            )
            self.mac = _Poly1305Mac(self.key)

        def test_constants(self):
            self.assertEqual(
                self.mac._s,
                0x1bf54941aff6bf4afdb20dfb8a800301
            )
            self.assertEqual(
                self.mac._r,
                0x806d5400e52447c036d555408bed685
            )
        
        
        def test_mac(self):
            _pt = b'Cryptographic Forum Research Group'
            pt = (_pt[i:i+16] for i in range(0, len(_pt), 16))
            
            self.mac.update(next(pt))
            self.assertEqual(
                self.mac._a, 
                0x2c88c77849d64ae9147ddeb88e69c83fc
            )
            
            self.mac.update(next(pt))
            self.assertEqual(
                self.mac._a, 
                0x2d8adaf23b0337fa7cccfb4ea344b30de
            )
            
            # object doesn't actually update on <16 bytes, so we do it manually
            a = self.mac._update(next(pt))
            self.assertEqual(
                a, 
                0x28d31b7caff946c77c8844335369d03a7
            )
            self.mac._a = 0x28d31b7caff946c77c8844335369d03a7
            
            self.assertEqual(
                self.mac.digest(),
                bytes.fromhex('a8061dc1305136c6c22b8baf0c0127a9')
            )

        
        
    class TestAEAD(unittest.TestCase):
        def setUp(self):
            self.key = bytes.fromhex(
                '80 81 82 83 84 85 86 87 88 89 8a 8b 8c 8d 8e 8f'
                '90 91 92 93 94 95 96 97 98 99 9a 9b 9c 9d 9e 9f'
            )
            self.iv = bytes.fromhex('07 00 00 00 40 41 42 43 44 45 46 47')
            self.auth = bytes.fromhex('50 51 52 53 c0 c1 c2 c3 c4 c5 c6 c7')
            self.plaintext = (
                b"Ladies and Gentlemen of the class of '99: "
                b'If I could offer you only one tip for the '
                b'future, sunscreen would be it.'
            )
            self.cipher = ChaCha20Poly1305Aead(self.key)
            
        
        def test_encryption_decryption(self):
            e = self.cipher.encryptor(self.iv, auth=self.auth)
            d = self.cipher.decryptor(self.iv, auth=self.auth)
            
            self.assertEqual(
                d.decrypt(e.encrypt(self.plaintext)),
                self.plaintext
            )
            
    
    class TestPFS(unittest.TestCase):
        def test_ratchet(self):
            pt = b'test'
            cipher = ChaCha()
            e = cipher.encryptor()
            d = cipher.decryptor(e.nonce)
            e.set_ratchet_friend(d)
            
            e.ratchet(b'123')
            self.assertEqual(
                d.decrypt(e.encrypt(pt)),
                pt
            )
            
            pt = b'A'*17 + b'B'*17
            
            e = cipher.encryptor()
            d = cipher.decryptor(e.nonce)
            e.set_ratchet_friend(d)
            
            plaintext = d.update(e.update(pt[:17]))
            self.assertEqual(
                plaintext, 
                pt[:17]
            )
            
            d.ratchet(b'123')  # test that e ratchets from d
            plaintext += d.update(e.update(pt[17:]))
            self.assertEqual(
                plaintext, 
                pt
            )
            
            
            
    unittest.main()