From 0c5afe44e7292559058a5f6eb2c76d274f3cd64e Mon Sep 17 00:00:00 2001 From: Vadim Lebedev Date: Mon, 18 Jul 2022 19:08:31 +0200 Subject: [PATCH] Fixes to handle password protected archives --- libarchive/__init__.py | 14 ++++++++++++-- libarchive/_libarchive.i | 5 +++-- libarchive/_libarchive_wrap.c | 5 +++-- libarchive/zip.py | 19 ++++++++++--------- 4 files changed, 28 insertions(+), 15 deletions(-) diff --git a/libarchive/__init__.py b/libarchive/__init__.py index c4861f1..9df1921 100644 --- a/libarchive/__init__.py +++ b/libarchive/__init__.py @@ -426,11 +426,13 @@ class Entry(object): class Archive(object): '''A low-level archive reader which provides forward-only iteration. Consider this a light-weight pythonic libarchive wrapper.''' - def __init__(self, f, mode='r', format=None, filter=None, entry_class=Entry, encoding=ENCODING, blocksize=BLOCK_SIZE): + def __init__(self, f, mode='r', format=None, filter=None, entry_class=Entry, + encoding=ENCODING, blocksize=BLOCK_SIZE, password=None): assert mode in ('r', 'w', 'wb', 'a'), 'Mode should be "r", "w", "wb", or "a".' self._stream = None self.encoding = encoding self.blocksize = blocksize + self.password = password if isinstance(f, str): self.filename = f f = open(f, mode) @@ -499,8 +501,12 @@ class Archive(object): self.format_func(self._a) self.filter_func(self._a) if self.mode == 'r': + if self.password: + self.add_passphrase(self.password) call_and_check(_libarchive.archive_read_open_fd, self._a, self._a, self.f.fileno(), self.blocksize) else: + if self.password: + self.set_passphrase(self.password) call_and_check(_libarchive.archive_write_open_fd, self._a, self._a, self.f.fileno()) def denit(self): @@ -562,7 +568,7 @@ class Archive(object): '''Write current archive entry contents to file. f can be a file-like object or a path.''' if isinstance(f, str): - basedir = os.path.basename(f) + basedir = os.path.dirname(f) if not os.path.exists(basedir): os.makedirs(basedir) f = open(f, 'w') @@ -626,6 +632,10 @@ class Archive(object): def add_passphrase(self, password): '''Adds a password to the archive.''' _libarchive.archive_read_add_passphrase(self._a, password) + + def set_passphrase(self, password): + '''Sets a password for the archive.''' + _libarchive.archive_write_set_passphrase(self._a, password) class SeekableArchive(Archive): diff --git a/libarchive/_libarchive.i b/libarchive/_libarchive.i index 56d29fa..6a09bd2 100644 --- a/libarchive/_libarchive.i +++ b/libarchive/_libarchive.i @@ -535,8 +535,9 @@ PyObject *archive_read_data_into_str(struct archive *archive, int len) { } PyObject *archive_write_data_from_str(struct archive *archive, PyObject *str) { - int len = PyString_Size(str); - if (!archive_write_data(archive, PyString_AS_STRING(str), len)) { + Py_ssize_t len = PyBytes_Size(str); + + if (!archive_write_data(archive, PyBytes_AS_STRING(str), len)) { PyErr_SetString(PyExc_RuntimeError, "could not write requested data."); return NULL; } diff --git a/libarchive/_libarchive_wrap.c b/libarchive/_libarchive_wrap.c index 315ecd2..1adbd15 100644 --- a/libarchive/_libarchive_wrap.c +++ b/libarchive/_libarchive_wrap.c @@ -3207,8 +3207,9 @@ PyObject *archive_read_data_into_str(struct archive *archive, int len) { } PyObject *archive_write_data_from_str(struct archive *archive, PyObject *str) { - int len = PyString_Size(str); - if (!archive_write_data(archive, PyString_AS_STRING(str), len)) { + Py_ssize_t len = PyBytes_Size(str); + + if (!archive_write_data(archive, PyBytes_AS_STRING(str), len)) { PyErr_SetString(PyExc_RuntimeError, "could not write requested data."); return NULL; } diff --git a/libarchive/zip.py b/libarchive/zip.py index 053d712..2995aea 100644 --- a/libarchive/zip.py +++ b/libarchive/zip.py @@ -62,8 +62,8 @@ class ZipEntry(Entry): class ZipFile(SeekableArchive): - def __init__(self, f, mode='r', compression=ZIP_DEFLATED, allowZip64=False): - super(ZipFile, self).__init__(f, mode=mode, format='zip', entry_class=ZipEntry, encoding='CP437') + def __init__(self, f, mode='r', compression=ZIP_DEFLATED, allowZip64=False, password=None): + super(ZipFile, self).__init__(f, mode=mode, format='zip', entry_class=ZipEntry, encoding='CP437', password=password) if mode == 'w' and compression == ZIP_STORED: # Disable compression for writing. _libarchive.archive_write_set_format_option(self.archive._a, "zip", "compression", "store") @@ -72,38 +72,39 @@ class ZipFile(SeekableArchive): getinfo = SeekableArchive.getentry def namelist(self): - return list(self.iterpaths) + return list(self.iterpaths()) def infolist(self): return list(self) def open(self, name, mode, pwd=None): - if pwd: - raise NotImplemented('Encryption not supported.') if mode == 'r': + if pwd: + self.add_passphrase(pwd) return self.readstream(name) else: return self.writestream(name) def extract(self, name, path=None, pwd=None): if pwd: - raise NotImplemented('Encryption not supported.') + self.add_passphrase(pwd) if not path: path = os.getcwd() return self.readpath(name, os.path.join(path, name)) def extractall(self, path, names=None, pwd=None): if pwd: - raise NotImplemented('Encryption not supported.') + self.add_passphrase(pwd) if not names: names = self.namelist() if names: + print(f"Extracting {names} files.") for name in names: self.extract(name, path) def read(self, name, pwd=None): if pwd: - raise NotImplemented('Encryption not supported.') + self.add_passphrase(pwd) return self.read(name) def writestr(self, member, data, compress_type=None): @@ -112,7 +113,7 @@ class ZipFile(SeekableArchive): return self.write(member, data) def setpassword(self, pwd): - raise NotImplemented('Encryption not supported.') + return self.set_passphrase(pwd) def testzip(self): raise NotImplemented()