Merge pull request #15 from EGA-archive/performance

Improving performance

crypt4gh/exceptions.py

@@ -8,8 +8,9 @@
 import logging
 import errno
 
-from nacl.exceptions import InvalidkeyError, BadSignatureError, CryptoError
-from cryptography.exceptions import InvalidTag
+from nacl.exceptions import (InvalidkeyError,
+                             BadSignatureError,
+                             CryptoError)
 
 LOG = logging.getLogger(__name__)
 
@@ -36,7 +37,7 @@ def exit_on_invalid_passphrase(func):
     def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
-        except (InvalidTag) as e:
+        except CryptoError as e:
             LOG.error('Exiting for %r', e)
             print('Invalid Key or Passphrase', file=sys.stderr)
             sys.exit(2)

crypt4gh/header.py

@@ -4,12 +4,14 @@
 import os
 import logging
 from itertools import chain
-from types import GeneratorType
+# from types import GeneratorType
 
-from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
-from nacl.bindings import crypto_kx_client_session_keys, crypto_kx_server_session_keys
+from nacl.bindings import (crypto_kx_client_session_keys,
+                           crypto_kx_server_session_keys,
+                           crypto_aead_chacha20poly1305_ietf_encrypt,
+                           crypto_aead_chacha20poly1305_ietf_decrypt)
+from nacl.exceptions import CryptoError
 from nacl.public import PrivateKey
-from cryptography.exceptions import InvalidTag
 
 from . import SEGMENT_SIZE, VERSION
 
@@ -182,11 +184,9 @@ def encrypt_X25519_Chacha20_Poly1305(data, seckey, recipient_pubkey):
     LOG.debug('shared key: %s', shared_key.hex())
 
     # Chacha20_Poly1305
-    engine = ChaCha20Poly1305(shared_key)
     nonce = os.urandom(12)
-    return (pubkey +
-            nonce +
-            engine.encrypt(nonce, data, None))  # No add
+    encrypted_data = crypto_aead_chacha20poly1305_ietf_encrypt(data, None, nonce, shared_key)  # no add
+    return (pubkey + nonce + encrypted_data)
 
 def decrypt_X25519_Chacha20_Poly1305(encrypted_part, privkey, sender_pubkey=None):
     #LOG.debug('-----------  Encrypted data: %s', encrypted_part.hex())
@@ -210,8 +210,7 @@ def decrypt_X25519_Chacha20_Poly1305(encrypted_part, privkey, sender_pubkey=None
     LOG.debug('shared key: %s', shared_key.hex())
 
     # Chacha20_Poly1305
-    engine = ChaCha20Poly1305(shared_key)
-    return engine.decrypt(nonce, packet_data, None)  # No add
+    return crypto_aead_chacha20poly1305_ietf_decrypt(packet_data, None, nonce, shared_key)  # no add
 
 
 def decrypt_packet(packet, keys, sender_pubkey=None):
@@ -233,7 +232,7 @@ def decrypt_packet(packet, keys, sender_pubkey=None):
             try:
                 privkey, _ = key # must fit
                 return decrypt_X25519_Chacha20_Poly1305(packet[4:], privkey, sender_pubkey=sender_pubkey)
-            except InvalidTag as tag:
+            except CryptoError as tag:
                 LOG.error('Packet Decryption failed: %s', tag)
             except Exception as e: # Any other error, like (IndexError, TypeError, ValueError)
                 LOG.error('Not a X25519 key: ignoring | %s', e)
@@ -305,8 +304,8 @@ def deconstruct(infile, keys, sender_pubkey=None):
 
     Leaves the infile stream right after the header.
 
-    :return: a pair with a list of ciphers and a generator of lengths from an edit list (or None if there was no edit list).
-    :rtype: (list of ChaCha20Poly1305 ciphers, int generator or None)
+    :return: a pair with a list of session keys and a generator of lengths from an edit list (or None if there was no edit list).
+    :rtype: (list of bytes, int generator or None)
 
     :raises: ValueError if the header could not be decrypted
     """
@@ -318,9 +317,9 @@ def deconstruct(infile, keys, sender_pubkey=None):
 
     data_packets, edit_packet = partition_packets(packets)
     # Parse returns the session key (since it should be method 0) 
-    ciphers = [ChaCha20Poly1305(parse_enc_packet(packet)) for packet in data_packets]
+    session_keys = [parse_enc_packet(packet) for packet in data_packets]
     edit_list = parse_edit_list_packet(edit_packet) if edit_packet else None
-    return ciphers, edit_list
+    return session_keys, edit_list
 
 # -------------------------------------
 # Header Re-Encryption

crypt4gh/keys/ssh.py

@@ -4,7 +4,7 @@
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives.ciphers import algorithms, Cipher, modes
-from cryptography.exceptions import InvalidTag
+from nacl.exceptions import CryptoError
 from nacl.bindings.crypto_sign import crypto_sign_ed25519_pk_to_curve25519, crypto_sign_ed25519_sk_to_curve25519
 
 from .kdf import derive_key
@@ -170,7 +170,7 @@ def parse_private_key(stream, callback):
 
     if private_data[:4] != private_data[4:8]: # check don't pass
         LOG.debug('Check: %s != %s', private_data[:4], private_data[4:8])
-        raise InvalidTag() 
+        raise CryptoError()
     private_data = io.BytesIO(private_data[8:])
     # Note: we ignore the comment and padding after the priv blob
     return _get_skpk_from_private_blob(private_data) # no need to unpad

crypt4gh/lib.py

@@ -6,11 +6,13 @@
 import io
 import collections
 
-from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
-from cryptography.exceptions import InvalidTag
+from nacl.bindings import (crypto_aead_chacha20poly1305_ietf_encrypt,
+                           crypto_aead_chacha20poly1305_ietf_decrypt)
+from nacl.exceptions import CryptoError
+
 
 from . import SEGMENT_SIZE
-from .exceptions import convert_error, close_on_broken_pipe
+from .exceptions import close_on_broken_pipe
 from . import header
 
 LOG = logging.getLogger(__name__)
@@ -32,13 +34,13 @@
 ##
 ##############################################################
 
-def _encrypt_segment(data, process, cipher):
+def _encrypt_segment(data, process, key):
     '''Utility function to generate a nonce, encrypt data with Chacha20, and authenticate it with Poly1305.'''
 
     #LOG.debug("Segment [%d bytes]: %s..%s", len(data), data[:10], data[-10:])
 
     nonce = os.urandom(12)
-    encrypted_data = cipher.encrypt(nonce, data, None)  # No add
+    encrypted_data = crypto_aead_chacha20poly1305_ietf_encrypt(data, None, nonce, key)  # no add
     process(nonce) # after producing the segment, so we don't start outputing when an error occurs
     process(encrypted_data)
 
@@ -81,7 +83,6 @@ def encrypt(keys, infile, outfile, offset=0, span=None):
     # Preparing the encryption engine
     encryption_method = 0 # only choice for this version
     session_key = os.urandom(32) # we use one session key for all blocks
-    cipher = ChaCha20Poly1305(session_key) # create a new one in case an old one is not reset
 
     # Output the header
     LOG.debug('Creating Crypt4GH header')
@@ -108,11 +109,11 @@ def encrypt(keys, infile, outfile, offset=0, span=None):
 
             if segment_len < SEGMENT_SIZE: # not a full segment
                 data = bytes(segment[:segment_len]) # to discard the bytes from the previous segments
-                _encrypt_segment(data, outfile.write, cipher)
+                _encrypt_segment(data, outfile.write, session_key)
                 break
 
             data = bytes(segment) # this is a full segment
-            _encrypt_segment(data, outfile.write, cipher)
+            _encrypt_segment(data, outfile.write, session_key)
 
     else: # we have a max size
         assert( span )
@@ -128,10 +129,10 @@ def encrypt(keys, infile, outfile, offset=0, span=None):
 
             if span < segment_len: # stop early
                 data = data[:span]
-                _encrypt_segment(data, outfile.write, cipher)
+                _encrypt_segment(data, outfile.write, session_key)
                 break
 
-            _encrypt_segment(data, outfile.write, cipher)
+            _encrypt_segment(data, outfile.write, session_key)
 
             span -= segment_len
 
@@ -156,17 +157,17 @@ def cipher_chunker(f, size):
         assert( ciphersegment_len > CIPHER_DIFF )
         yield ciphersegment
 
-def decrypt_block(ciphersegment, ciphers):
+def decrypt_block(ciphersegment, session_keys):
     # Trying the different session keys (via the cipher objects)
     # Note: we could order them and if one fails, we move it at the end of the list
     # So... LRU solution. For now, try them as they come.
     nonce = ciphersegment[:12]
     data = ciphersegment[12:]
 
-    for cipher in ciphers:
+    for key in session_keys:
         try:
-            return cipher.decrypt(nonce, data, None)  # No aad, and break the loop
-        except InvalidTag as tag:
+            return crypto_aead_chacha20poly1305_ietf_decrypt(data, None, nonce, key)  # no add, and break the loop
+        except CryptoError as tag:
             LOG.error('Decryption failed: %s', tag)
     else: # no cipher worked: Bark!
         raise ValueError('Could not decrypt that block')
@@ -214,7 +215,7 @@ def limited_output(offset=0, limit=None, process=None):
         offset = 0 # reset offset
 
 
-def body_decrypt(infile, ciphers, output, offset):
+def body_decrypt(infile, session_keys, output, offset):
     """Decrypt the whole data portion.
 
     We fast-forward if offset >= SEGMENT_SIZE.
@@ -231,16 +232,16 @@ def body_decrypt(infile, ciphers, output, offset):
 
     try:
         for ciphersegment in cipher_chunker(infile, CIPHER_SEGMENT_SIZE):
-            segment = decrypt_block(ciphersegment, ciphers)
+            segment = decrypt_block(ciphersegment, session_keys)
             output.send(segment)
     except ProcessingOver:  # output raised it
         pass
 
 
 class DecryptedBuffer():
-    def __init__(self, fileobj, ciphers, output):
+    def __init__(self, fileobj, session_keys, output):
         self.fileobj = fileobj
-        self.ciphers = ciphers
+        self.session_keys = session_keys
         self.buf = io.BytesIO()
         self.block = 0  # just used for printing, if that block is entirely skipped
         self.output = output
@@ -278,7 +279,7 @@ def _fetch(self, nodecrypt=False):
         # else, we decrypt
         LOG.debug('Decrypting block %d', self.block)
         assert( len(data) > CIPHER_DIFF )
-        segment = decrypt_block(data, self.ciphers)
+        segment = decrypt_block(data, self.session_keys)
         LOG.debug('Adding %d bytes to the buffer', len(segment))
         self._append_to_buf(segment)
         LOG.debug('Buffer size: %d', self.buf_size())
@@ -322,7 +323,7 @@ def read(self, size):
             size -= len(b2)
 
 
-def body_decrypt_parts(infile, ciphers, output, edit_list=None):
+def body_decrypt_parts(infile, session_keys, output, edit_list=None):
     """Decrypt the data portion according to the edit list.
 
     We do not decrypt segments that are entirely skipped, and only output a warning (that it should not be the case).
@@ -332,7 +333,7 @@ def body_decrypt_parts(infile, ciphers, output, edit_list=None):
     LOG.debug('Edit List: %s', edit_list)
     assert(len(edit_list) > 0), "You can not call this function without an edit_list"
 
-    decrypted = DecryptedBuffer(infile, ciphers, output)
+    decrypted = DecryptedBuffer(infile, session_keys, output)
 
     try:
 
@@ -374,7 +375,7 @@ def decrypt(keys, infile, outfile, sender_pubkey=None, offset=0, span=None):
         )
     )
 
-    ciphers, edit_list = header.deconstruct(infile, keys, sender_pubkey=sender_pubkey)
+    session_keys, edit_list = header.deconstruct(infile, keys, sender_pubkey=sender_pubkey)
 
     # Infile in now positioned at the beginning of the data portion
 
@@ -384,11 +385,11 @@ def decrypt(keys, infile, outfile, sender_pubkey=None, offset=0, span=None):
 
     if edit_list is None:
         # No edit list: decrypt all segments until the end
-        body_decrypt(infile, ciphers, output, offset)
+        body_decrypt(infile, session_keys, output, offset)
         # We could use body_decrypt_parts but there is an inner buffer, and segments might not be aligned
     else:
         # Edit list: it drives which segments is decrypted
-        body_decrypt_parts(infile, ciphers, output, edit_list=list(edit_list))
+        body_decrypt_parts(infile, session_keys, output, edit_list=list(edit_list))
 
     LOG.info('Decryption Over')
 

tests/_common/edit_list_gen.py

@@ -12,10 +12,8 @@
 from functools import partial
 from getpass import getpass
 
-from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
-
 from crypt4gh.keys import get_private_key, get_public_key
-from crypt4gh import header,lib, SEGMENT_SIZE
+from crypt4gh import header, lib, SEGMENT_SIZE
 
 if __name__ == '__main__':
 
@@ -73,7 +71,6 @@
     #############################################################
     encryption_method = 0 # only choice for this version
     session_key = os.urandom(32) # we use one session key for all blocks
-    cipher = ChaCha20Poly1305(session_key) # create a new one in case an old one is not reset
 
     #############################################################
     # Output the header
@@ -99,10 +96,10 @@
 
         if segment_len < SEGMENT_SIZE: # not a full segment
             data = bytes(segment[:segment_len]) # to discard the bytes from the previous segments
-            lib._encrypt_segment(data, outfile.write, cipher)
+            lib._encrypt_segment(data, outfile.write, session_key)
             break
 
         data = bytes(segment) # this is a full segment
-        lib._encrypt_segment(data, outfile.write, cipher)
+        lib._encrypt_segment(data, outfile.write, session_key)