Merge pull request #27 from isi-mfurer/integ-0.2.5

Integrate Version 0.2.5

pike/core.py

@@ -697,3 +697,44 @@ def __or__(self, o):
 
     def __and__(self, o):
         return self.__class__(super(FlagEnum,self).__and__(o))
+
+class Let(object):
+    """
+    A Let allows for temporarily storing passed arguments in the _settings
+    member of a target object for the duration of the contextmanager's scope
+
+    Implementation on a factory class
+
+    class MyThingFactory(object):
+        def __init__(self):
+            self._settings = {}
+        def generate_thing(self):
+            # do something here
+            new_obj = Thing()
+            for k,v in self._settings.iteritems():
+                setattr(new_obj, k, v)
+            return new_obj
+        def let(self, **kwds):
+            return Let(self, kwds)
+
+    Calling code that wants to temporarily customize the Thing objects that are
+    returned by the factory:
+
+    fac = MyThingFactory()
+    with fac.let(this="that", foo="bar"):
+        a_thing = fac.generate_thing()
+        self.assertEqual(a_thing.this, "that")
+    """
+    def __init__(self, target, settings_dict):
+        self.target = target
+        self.settings_dict = settings_dict
+        if not hasattr(target, "_settings"):
+            target._settings = {}
+
+    def __enter__(self):
+        self.backup = dict(self.target._settings)
+        self.target._settings.update(self.settings_dict)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.target._settings = self.backup

pike/crypto.py

@@ -42,25 +42,49 @@
 
 from Cryptodome.Cipher import AES
 
+
+def pad_right(value, length, byte='\0'):
+    if len(value) > length:
+        value = value[:length]
+    elif len(value) < 16:
+        value += array.array('B', byte*(length - len(value)))
+    return value
+
+
+class CipherMismatch(Exception):
+    pass
+
+
 class Ciphers(core.ValueEnum):
+    SMB2_NONE_CIPHER = 0x0000
     SMB2_AES_128_CCM = 0x0001
     SMB2_AES_128_GCM = 0x0002
 
 Ciphers.import_items(globals())
 
+cipher_map = {
+        SMB2_AES_128_CCM: (AES.MODE_CCM, 11),
+        SMB2_AES_128_GCM: (AES.MODE_GCM, 12)
+}
+
+
 class EncryptionCapabilities(core.Frame):
     context_type = smb2.SMB2_ENCRYPTION_CAPABILITIES
+
     def __init__(self):
         self.ciphers = []
+        self.ciphers_count = None
 
     def _encode(self, cur):
-        cur.encode_uint16le(len(self.ciphers))
+        if self.ciphers_count is None:
+            self.ciphers_count = len(self.ciphers)
+        cur.encode_uint16le(self.ciphers_count)
         for c in self.ciphers:
             cur.encode_uint16le(c)
 
     def _decode(self, cur):
-        cipher_count = cur.decode_uint16le()
-        for ix in xrange(cipher_count):
+        self.ciphers_count = cur.decode_uint16le()
+        for ix in xrange(self.ciphers_count):
             self.ciphers.append(Ciphers(cur.decode_uint16le()))
 
 
@@ -94,12 +118,17 @@ def __init__(self, parent):
         core.Frame.__init__(self, parent)
         self.protocol_id = array.array('B', "\xfdSMB")
         self.signature = None
+        # the value of nonce is always used in the encryption routine
         self.nonce = array.array('B',
                                  map(random.randint, [0]*16, [255]*16))
-        self.original_message_size = 0
+        # if wire_nonce is set, it will be sent on the wire instead of nonce
+        self.wire_nonce = None
+        self.original_message_size = None
+        self.reserved = None
         self.flags = 0x1
         self.session_id = None
         self.encryption_context = None
+        self.additional_authenticated_data_buf = array.array('B')
         if parent is not None:
             parent.transform = self
         else:
@@ -130,42 +159,56 @@ def _encode_header(self, cur):
         # the signature will be written in _encode_smb2
         self.signature_offset = cur.offset
         cur.advanceto(cur + 16)
+        # the crypto header will be written in _encode_smb2
+        self.crypto_header_offset = cur.offset
+        # save a space for the wire nonce
+        self.wire_nonce_hole = cur.hole.encode_bytes('\0'*16)
+        cur.advanceto(cur + 16)
 
         # the following fields are part of AdditionalAuthenticatedData and are
         # used as inputs to the AES cipher
-        self.crypto_header_start = cur.offset
+        aad_cur = core.Cursor(self.additional_authenticated_data_buf, 0)
         # nonce field size is 16 bytes right padded with zeros
-        if len(self.nonce) > 16:
-            self.nonce = self.nonce[:16]
-        elif len(self.nonce) < 16:
-            self.nonce = self.nonce + array.array('B', "\0"*(16-len(self.nonce)))
-        cur.encode_bytes(self.nonce)
-        self.original_message_size_hole = cur.hole.encode_uint32le(0)
-        cur.encode_uint16le(0)  # reserved
-        cur.encode_uint16le(self.flags)
-        cur.encode_uint64le(self.session_id)
-        self.crypto_header_end = cur.offset
+        self.nonce = pad_right(self.nonce, 16)
+        aad_cur.encode_bytes(self.nonce)
+        self.original_message_size_hole = aad_cur.hole.encode_uint32le(0)
+        if self.reserved is None:
+            self.reserved = 0
+        aad_cur.encode_uint16le(self.reserved)  # reserved
+        aad_cur.encode_uint16le(self.flags)
+        aad_cur.encode_uint64le(self.session_id)
 
     def _encode_smb2(self, cur):
         # serialize all chained Smb2 commands into one buffer
         original_message_buf = array.array('B')
         original_message_cur = core.Cursor(original_message_buf, 0)
         for smb_frame in self.parent:
             smb_frame.encode(original_message_cur)
-        self.original_message_size = len(original_message_buf)
+        if self.original_message_size is None:
+            self.original_message_size = len(original_message_buf)
         self.original_message_size_hole(self.original_message_size)
-        crypto_header = cur.array[self.crypto_header_start:self.crypto_header_end]
         (self.ciphertext,
-         self.signature) = self.encryption_context.encrypt(
+         crypto_hmac) = self.encryption_context.encrypt(
                                 original_message_buf,
-                                crypto_header,
+                                self.additional_authenticated_data_buf,
                                 self.nonce)
         cur.encode_bytes(self.ciphertext)
 
         # fill in the signature hole
         sig_cur = core.Cursor(cur.array, self.signature_offset)
+        if self.signature is None:
+            self.signature = crypto_hmac
         sig_cur.encode_bytes(self.signature)
 
+        # fill in the header
+        aad_cur = core.Cursor(cur.array, self.crypto_header_offset)
+        aad_cur.encode_bytes(self.additional_authenticated_data_buf)
+
+        # fill in the wire nonce
+        if self.wire_nonce is None:
+            self.wire_nonce = self.nonce
+        self.wire_nonce_hole(pad_right(self.wire_nonce, 16))
+
     def _decode(self, cur):
         self._decode_header(cur)
         self._decode_smb2(cur)
@@ -178,7 +221,7 @@ def _decode_header(self, cur):
         self.crypto_header_start = cur.offset
         self.nonce = cur.decode_bytes(16)
         self.original_message_size = cur.decode_uint32le()
-        cur.decode_uint16le()   # reserved
+        self.reserved = cur.decode_uint16le()   # reserved
         self.flags = cur.decode_uint16le()
         self.session_id = cur.decode_uint64le()
         self.crypto_header_end = cur.offset
@@ -232,13 +275,15 @@ class EncryptionContext(object):
     """
     def __init__(self, keys, ciphers):
         self.keys = keys
-        self.cipher = ciphers[0]
-        if self.cipher == SMB2_AES_128_CCM:
-            self.aes_mode = AES.MODE_CCM
-            self.nonce_length = 11
-        elif self.cipher == SMB2_AES_128_GCM:
-            self.aes_mode = AES.MODE_GCM
-            self.nonce_length = 12
+        for c in ciphers:
+            if c in cipher_map:
+                self.aes_mode, self.nonce_length = cipher_map[c]
+                self.cipher = c
+                break
+        else:
+            raise CipherMismatch(
+                    "Client did not recognize any ciphers returned by the "
+                    "server: {0}".format(ciphers))
 
     def encrypt(self, plaintext, authenticated_data, nonce):
         enc_cipher = AES.new(self.keys.encryption,
@@ -253,6 +298,8 @@ def decrypt(self, ciphertext, signature, authenticated_data, nonce):
                              self.aes_mode,
                              nonce=nonce[:self.nonce_length].tostring())
         dec_cipher.update(authenticated_data.tostring())
-        return array.array('B',
-                dec_cipher.decrypt_and_verify(ciphertext.tostring(),
-                                              signature.tostring()))
+        return array.array(
+                'B',
+                dec_cipher.decrypt_and_verify(
+                    ciphertext.tostring(),
+                    signature.tostring()))