express: Add support for bytes multifile encryption passwords (#1334)

panda/src/express/CMakeLists.txt

@@ -141,6 +141,8 @@ set(P3EXPRESS_IGATEEXT
   virtualFileSystem_ext.h
   virtualFile_ext.cxx
   virtualFile_ext.h
+  multifile_ext.h
+  multifile_ext.I
 )
 
 composite_sources(p3express P3EXPRESS_SOURCES)

panda/src/express/multifile.h

@@ -69,8 +69,6 @@ PUBLISHED:
 
   INLINE void set_encryption_flag(bool flag);
   INLINE bool get_encryption_flag() const;
-  INLINE void set_encryption_password(const std::string &encryption_password);
-  INLINE const std::string &get_encryption_password() const;
 
   INLINE void set_encryption_algorithm(const std::string &encryption_algorithm);
   INLINE const std::string &get_encryption_algorithm() const;
@@ -86,6 +84,9 @@ PUBLISHED:
   std::string update_subfile(const std::string &subfile_name, const Filename &filename,
                         int compression_level);
 
+  EXTENSION(INLINE PyObject *set_encryption_password(PyObject *encryption_password) const);
+  EXTENSION(INLINE PyObject *get_encryption_password() const);
+
 #ifdef HAVE_OPENSSL
   bool add_signature(const Filename &certificate,
                      const Filename &chain,
@@ -143,6 +144,9 @@ PUBLISHED:
   INLINE const std::string &get_header_prefix() const;
 
 public:
+  INLINE void set_encryption_password(const std::string &encryption_password);
+  INLINE const std::string &get_encryption_password() const;
+
 #ifdef HAVE_OPENSSL
   class CertRecord {
   public:

panda/src/express/multifile_ext.I

@@ -0,0 +1,79 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file multifile_ext.I
+ * @author Derzsi Daniel
+ * @date 2022-07-20
+ */
+
+#include <cstring>
+
+/**
+ * Specifies the password, either as a Python string or a Python bytes object,
+ * that will be used to encrypt subfiles subsequently added to the multifile
+ */
+INLINE PyObject *Extension<Multifile>::
+set_encryption_password(PyObject *encryption_password) const {
+  Py_ssize_t pass_len;
+
+  // Have we been passed a string?
+  if (PyUnicode_Check(encryption_password)) {
+    const char *pass_str = PyUnicode_AsUTF8AndSize(encryption_password, &pass_len);
+    _this->set_encryption_password(std::string(pass_str, pass_len));
+    return Dtool_Return_None();
+  }
+
+  // Have we been passed a bytes object?
+  if (PyBytes_Check(encryption_password)) {
+    char *pass_str;
+
+    if (PyBytes_AsStringAndSize(encryption_password, &pass_str, &pass_len) < 0) {
+      PyErr_SetString(PyExc_TypeError, "A valid bytes object is required.");
+      return NULL;
+    }
+
+    // It is dangerous to use null bytes inside the encryption password.
+    // OpenSSL will cut off the password prematurely at the first null byte
+    // encountered.
+    if (memchr(pass_str, '\0', pass_len) != NULL) {
+      PyErr_SetString(PyExc_ValueError, "The password must not contain null bytes.");
+      return NULL;
+    }
+
+    _this->set_encryption_password(std::string(pass_str, pass_len));
+    return Dtool_Return_None();
+  }
+
+  return Dtool_Raise_BadArgumentsError(
+    "set_encryption_password(const Multifile self, str encryption_password)\n"
+  );
+}
+
+/**
+ * Returns the password that will be used to encrypt subfiles subsequently
+ * added to the multifile, either as a Python string (when possible) or a
+ * Python bytes object.
+ */
+INLINE PyObject *Extension<Multifile>::
+get_encryption_password() const {
+  std::string password = _this->get_encryption_password();
+  const char *pass_str = password.c_str();
+  Py_ssize_t pass_len = password.length();
+
+  // First, attempt to decode it as an UTF-8 string...
+  PyObject *result = PyUnicode_DecodeUTF8(pass_str, pass_len, NULL);
+
+  if (PyErr_Occurred()) {
+    // This password cannot be decoded as an UTF-8 string, so let's
+    // return it as a bytes object.
+    PyErr_Clear();
+    result = PyBytes_FromStringAndSize(pass_str, pass_len);
+  }
+
+  return result;
+}

panda/src/express/multifile_ext.h

@@ -0,0 +1,40 @@
+/**
+ * PANDA 3D SOFTWARE
+ * Copyright (c) Carnegie Mellon University.  All rights reserved.
+ *
+ * All use of this software is subject to the terms of the revised BSD
+ * license.  You should have received a copy of this license along
+ * with this source code in a file named "LICENSE."
+ *
+ * @file multifile_ext.h
+ * @author Derzsi Daniel
+ * @date 2022-07-20
+ */
+
+#ifndef MULTIFILE_EXT_H
+#define MULTIFILE_EXT_H
+
+#include "dtoolbase.h"
+
+#ifdef HAVE_PYTHON
+
+#include "extension.h"
+#include "multifile.h"
+#include "py_panda.h"
+
+/**
+ * This class defines the extension methods for Multifile, which are called
+ * instead of any C++ methods with the same prototype.
+ */
+template<>
+class Extension<Multifile> : public ExtensionBase<Multifile> {
+public:
+  INLINE PyObject *set_encryption_password(PyObject *encryption_password) const;
+  INLINE PyObject *get_encryption_password() const;
+};
+
+#include "multifile_ext.I"
+
+#endif  // HAVE_PYTHON
+
+#endif  // MULTIFILE_EXT_H

panda/src/express/p3express_ext_composite.cxx

@@ -3,3 +3,4 @@
 #include "stringStream_ext.cxx"
 #include "virtualFileSystem_ext.cxx"
 #include "virtualFile_ext.cxx"
+#include "multifile_ext.h"

tests/express/test_multifile.py

@@ -10,3 +10,16 @@ def test_multifile_read_empty():
     assert m.is_read_valid()
     assert m.get_num_subfiles() == 0
     m.close()
+
+
+def test_multifile_password():
+    m = Multifile()
+
+    m.set_encryption_password('Panda3D rocks!')
+    assert m.get_encryption_password() == 'Panda3D rocks!'
+
+    m.set_encryption_password(b'Panda3D is awesome!')
+    assert m.get_encryption_password() == 'Panda3D is awesome!'
+
+    m.set_encryption_password(b'\xc4\x97\xa1\x01\x85\xb6')
+    assert m.get_encryption_password() == b'\xc4\x97\xa1\x01\x85\xb6'