added initial support for AES256-SHA256 encrpyted archives (Bugzilla #7)

py7zlib.py

@@ -37,6 +37,19 @@
     # reduce is available in functools starting with Python 2.6
     pass
 
+try:
+    import M2Crypto
+    from M2Crypto import EVP
+except ImportError:
+    # support for encrypted files is optional
+    M2Crypto = None
+
+try:
+    from hashlib import sha256
+except ImportError:
+    # hashlib is optional for some encrypted archives
+    sha256 = None
+
 READ_BLOCKSIZE                   = 16384
 
 MAGIC_7Z                         = '7z\xbc\xaf\x27\x1c'
@@ -72,6 +85,7 @@
 COMPRESSION_METHOD_MISC          = '\x04'
 COMPRESSION_METHOD_MISC_ZIP      = '\x04\x01'
 COMPRESSION_METHOD_MISC_BZIP     = '\x04\x02'
+COMPRESSION_METHOD_7Z_AES256_SHA256 = '\x06\xf1\x07\x01'
 
 class ArchiveError(Exception):
     pass
@@ -85,6 +99,15 @@ class EncryptedArchiveError(ArchiveError):
 class UnsupportedCompressionMethodError(ArchiveError):
     pass
 
+class DecryptionError(ArchiveError):
+    pass
+
+class NoPasswordGivenError(DecryptionError):
+    pass
+
+class WrongPasswordError(DecryptionError):
+    pass
+
 class Base(object):
     """ base class with support for various basic read/write functions """
 
@@ -461,6 +484,14 @@ def __init__(self, info, start, src_start, size, folder, archive, maxsize=None):
             COMPRESSION_METHOD_MISC_ZIP: '_read_zip',
             COMPRESSION_METHOD_MISC_BZIP: '_read_bzip',
         }
+        if M2Crypto is not None:
+            if sha256 is not None:
+                self._decoders.update({
+                    COMPRESSION_METHOD_7Z_AES256_SHA256: '_read_7z_aes256_sha256',
+                })
+
+    def _is_encrypted(self):
+        return COMPRESSION_METHOD_7Z_AES256_SHA256 in [x['method'] for x in self._folder.coders]
 
     def reset(self):
         self.pos = 0
@@ -472,26 +503,25 @@ def read(self):
         data = None
         for coder in self._folder.coders:
             method = coder['method']
-            if method[0] == COMPRESSION_METHOD_CRYPTO:
-                raise EncryptedArchiveError("encrypted archives are not supported yet")
-
             decoder = None
             while method and decoder is None:
                 decoder = self._decoders.get(method, None)
                 method = method[:-1]
 
             if decoder is None:
-                raise UnsupportedCompressionMethodError(coder['method'])
+                raise UnsupportedCompressionMethodError(repr(coder['method']))
 
             data = getattr(self, decoder)(coder, data)
 
         return data
 
     def _read_copy(self, coder, input):
-        self._file.seek(self._src_start)
-        return self._file.read(self.uncompressed)[self._start:self._start+self.size]
+        if not input:
+            self._file.seek(self._src_start)
+            input = self._file.read(self.uncompressed)
+        return input[self._start:self._start+self.size]
 
-    def _read_from_decompressor(self, coder, decompressor, checkremaining=False):
+    def _read_from_decompressor(self, coder, decompressor, input, checkremaining=False):
         data = ''
         idx = 0
         cnt = 0
@@ -500,7 +530,7 @@ def _read_from_decompressor(self, coder, decompressor, checkremaining=False):
         if properties:
             decompressor.decompress(properties)
         total = self.compressed
-        if total is None:
+        if not input and total is None:
             remaining = self._start+self.size
             out = StringIO()
             while remaining > 0:
@@ -514,24 +544,71 @@ def _read_from_decompressor(self, coder, decompressor, checkremaining=False):
 
             data = out.getvalue()
         else:
+            if not input:
+                input = self._file.read(total)
             if checkremaining:
-                data = decompressor.decompress(self._file.read(total), self._start+self.size)
+                data = decompressor.decompress(input, self._start+self.size)
             else:
-                data = decompressor.decompress(self._file.read(total))
+                data = decompressor.decompress(input)
         return data[self._start:self._start+self.size]
 
     def _read_lzma(self, coder, input):
         dec = pylzma.decompressobj(maxlength=self._start+self.size)
-        return self._read_from_decompressor(coder, dec, checkremaining=True)
+        try:
+            return self._read_from_decompressor(coder, dec, input, checkremaining=True)
+        except ValueError:
+            if self._is_encrypted():
+                raise WrongPasswordError('invalid password')
+
+            raise
 
     def _read_zip(self, coder, input):
         dec = zlib.decompressobj(-15)
-        return self._read_from_decompressor(coder, dec, checkremaining=True)
+        return self._read_from_decompressor(coder, dec, input, checkremaining=True)
 
     def _read_bzip(self, coder, input):
         dec = bz2.BZ2Decompressor()
-        return self._read_from_decompressor(coder, dec)
+        return self._read_from_decompressor(coder, dec, input)
+
+    def _read_7z_aes256_sha256(self, coder, input):
+        if not self._archive.password:
+            raise NoPasswordGivenError()
+
+        # TODO: this needs some sanity checks
+        firstbyte = ord(coder['properties'][0])
+        numcyclespower = firstbyte & 0x3f
+        if firstbyte & 0xc0 != 0:
+            saltsize = (firstbyte >> 7) & 1
+            ivsize = (firstbyte >> 6) & 1
+
+            secondbyte = ord(coder['properties'][1])
+            saltsize += (secondbyte >> 4)
+            ivsize += (secondbyte & 0x0f)
+
+            assert len(coder['properties']) == 2+saltsize+ivsize
+            salt = coder['properties'][2:2+saltsize]
+            iv = coder['properties'][2+saltsize:2+saltsize+ivsize]
+            assert len(salt) == saltsize
+            assert len(iv) == ivsize
+            assert numcyclespower <= 24
+            if ivsize < 16:
+                iv += '\x00'*(16-ivsize)
+        else:
+            salt = iv = ''
 
+        password = self._archive.password.encode('utf-16-le')
+        key = pylzma.calculate_key(password, numcyclespower, salt=salt)
+        cipher = pylzma.AESDecrypt(key, iv=iv)
+        if not input:
+            self._file.seek(self._src_start)
+            uncompressed_size = self.uncompressed
+            if uncompressed_size & 0x0f:
+                # we need a multiple of 16 bytes
+                uncompressed_size += 16 - (uncompressed_size & 0x0f)
+            input = self._file.read(uncompressed_size)
+        result = cipher.decrypt(input)
+        return result
+
     def checkcrc(self):
         if self.digest is None:
             return True
@@ -544,8 +621,9 @@ def checkcrc(self):
 class Archive7z(Base):
     """ the archive itself """
 
-    def __init__(self, file):
+    def __init__(self, file, password=None):
         self._file = file
+        self.password = password
         self.header = file.read(len(MAGIC_7Z))
         if self.header != MAGIC_7Z:
             raise FormatError, 'not a 7z file'

setup.py

@@ -53,7 +53,9 @@ class UnsupportedPlatformWarning(Warning):
 if IS_WINDOWS:
     libraries += ['user32', 'oleaut32']
 
-include_dirs = []
+include_dirs = [
+    'src/sdk',
+]
 
 if sys.platform == 'darwin':
     # additional include directories are required when compiling on Darwin platforms
@@ -110,11 +112,14 @@ def build_extension(self, ext):
 except: version = 'unknown'
 modules = ['py7zlib']
 c_files = ['src/pylzma/pylzma.c', 'src/pylzma/pylzma_decompressobj.c', 'src/pylzma/pylzma_compressfile.c',
-           'src/pylzma/pylzma_decompress.c', 'src/pylzma/pylzma_compress.c', 'src/pylzma/pylzma_streams.c']
+           'src/pylzma/pylzma_decompress.c', 'src/pylzma/pylzma_compress.c', 'src/pylzma/pylzma_streams.c', \
+           'src/pylzma/pylzma_aes.c']
 compile_args = []
 link_args = []
 macros = []
-lzma_files = ('src/sdk/LzFind.c', 'src/sdk/LzmaDec.c', 'src/sdk/LzmaEnc.c', )
+lzma_files = ('src/sdk/LzFind.c', 'src/sdk/LzmaDec.c', 'src/sdk/LzmaEnc.c', \
+    'src/7zip/C/CpuArch.c', 'src/7zip/C/Aes.c', 'src/7zip/C/AesOpt.c', \
+    'src/7zip/C/Sha256.c')
 if ENABLE_COMPATIBILITY:
     c_files += ('src/pylzma/pylzma_decompress_compat.c', 'src/pylzma/pylzma_decompressobj_compat.c', )
     lzma_files += ('src/compat/LzmaCompatDecode.c', )
@@ -151,6 +156,10 @@ def build_extension(self, ext):
     cmdclass = {
         'build_ext': build_ext,
     },
+    extras_require = {
+        'decrypt': ['m2crypto'],
+    },
+    tests_require = ['pylzma[decrypt]'],
     test_suite = 'tests.suite',
     zip_safe = False,
 )

src/pylzma/pylzma.c

@@ -27,6 +27,8 @@
 #include <cStringIO.h>
 
 #include "../sdk/7zVersion.h"
+#include "../7zip/C/Sha256.h"
+#include "../7zip/C/Aes.h"
 
 #include "pylzma.h"
 #include "pylzma_compress.h"
@@ -36,6 +38,7 @@
 #include "pylzma_compressobj.h"
 #endif
 #include "pylzma_compressfile.h"
+#include "pylzma_aes.h"
 #ifdef WITH_COMPAT
 #include "pylzma_decompress_compat.h"
 #include "pylzma_decompressobj_compat.h"
@@ -45,6 +48,79 @@
 PyInterpreterState* _pylzma_interpreterState = NULL;
 #endif
 
+const char
+doc_calculate_key[] = \
+    "calculate_key(password, cycles, salt=None, digest='sha256') -- Calculate decryption key.";
+
+static PyObject *
+pylzma_calculate_key(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    char *password;
+    int pwlen;
+    int cycles;
+    PyObject *pysalt=NULL;
+    char *salt;
+    int saltlen;
+    char *digest="sha256";
+    char key[32];
+    // possible keywords for this function
+    static char *kwlist[] = {"password", "cycles", "salt", "digest", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s#i|Os", kwlist, &password, &pwlen, &cycles, &pysalt, &digest))
+        return NULL;
+
+    if (pysalt == Py_None) {
+        pysalt = NULL;
+    } else if (!PyString_Check(pysalt)) {
+        PyErr_Format(PyExc_TypeError, "salt must be a string, got a %s", pysalt->ob_type->tp_name);
+        return NULL;
+    }
+
+    if (strcmp(digest, "sha256") != 0) {
+        PyErr_Format(PyExc_TypeError, "digest %s is unsupported", digest);
+        return NULL;
+    }
+
+    if (pysalt != NULL) {
+        salt = PyString_AS_STRING(pysalt);
+        saltlen = PyString_Size(pysalt);
+    } else {
+        salt = NULL;
+        saltlen = 0;
+    }
+
+    if (cycles == 0x3f) {
+        int pos;
+        int i;
+        for (pos = 0; pos < saltlen; pos++)
+            key[pos] = salt[pos];
+        for (i = 0; i<pwlen && pos < 32; i++)
+            key[pos++] = password[i];
+        for (; pos < 32; pos++)
+            key[pos] = 0;
+    } else {
+        Py_BEGIN_ALLOW_THREADS
+        CSha256 sha;
+        Sha256_Init(&sha);
+        long round;
+        int i;
+        long rounds = (long) 1 << cycles;
+        unsigned char temp[8] = { 0,0,0,0,0,0,0,0 };
+        for (round = 0; round < rounds; round++) {
+            Sha256_Update(&sha, (Byte *) salt, saltlen);
+            Sha256_Update(&sha, (Byte *) password, pwlen);
+            Sha256_Update(&sha, (Byte *) temp, 8);
+            for (i = 0; i < 8; i++)
+                if (++(temp[i]) != 0)
+                    break;
+        }
+        Sha256_Final(&sha, (Byte *) &key);
+        Py_END_ALLOW_THREADS
+    }
+
+    return PyString_FromStringAndSize(key, 32);
+}
+
 PyMethodDef
 methods[] = {
     // exported functions
@@ -55,6 +131,7 @@ methods[] = {
     {"decompress_compat",    (PyCFunction)pylzma_decompress_compat,    METH_VARARGS | METH_KEYWORDS, (char *)&doc_decompress_compat},
     {"decompressobj_compat", (PyCFunction)pylzma_decompressobj_compat, METH_VARARGS,                 (char *)&doc_decompressobj_compat},
 #endif
+    {"calculate_key",   (PyCFunction)pylzma_calculate_key,  METH_VARARGS | METH_KEYWORDS,   (char *)&doc_calculate_key},
     {NULL, NULL},
 };
 
@@ -79,6 +156,10 @@ initpylzma(void)
     if (PyType_Ready(&CCompressionFileObject_Type) < 0)
         return;
 
+    CAESDecrypt_Type.tp_new = PyType_GenericNew;
+    if (PyType_Ready(&CAESDecrypt_Type) < 0)
+        return;
+
     m = Py_InitModule("pylzma", methods);
 
     Py_INCREF(&CDecompressionObject_Type);
@@ -90,6 +171,9 @@ initpylzma(void)
     Py_INCREF(&CCompressionFileObject_Type);
     PyModule_AddObject(m, "compressfile", (PyObject *)&CCompressionFileObject_Type);
 
+    Py_INCREF(&CAESDecrypt_Type);
+    PyModule_AddObject(m, "AESDecrypt", (PyObject *)&CAESDecrypt_Type);
+
     PyModule_AddIntConstant(m, "SDK_VER_MAJOR", MY_VER_MAJOR);
     PyModule_AddIntConstant(m, "SDK_VER_MINOR", MY_VER_MINOR);
     PyModule_AddIntConstant(m, "SDK_VER_BUILD ", MY_VER_BUILD);
@@ -98,6 +182,7 @@ initpylzma(void)
     PyModule_AddStringConstant(m, "SDK_COPYRIGHT", MY_COPYRIGHT);
     PyModule_AddStringConstant(m, "SDK_VERSION_COPYRIGHT_DATE", MY_VERSION_COPYRIGHT_DATE);
 
+    AesGenTables();
     PycString_IMPORT;
 
 #if defined(WITH_THREAD)