Merge pull request nolze#59 from nolze/feature/exceptions

Use custom exceptions

msoffcrypto/__init__.py

@@ -2,6 +2,8 @@
 import zipfile
 import pkg_resources
 
+from . import exceptions
+
 __version__ = pkg_resources.get_distribution("msoffcrypto-tool").version
 
 
@@ -20,6 +22,23 @@ def OfficeFile(file):
  ... officefile.keyTypes
  ('password', 'private_key', 'secret_key')
 
+ >>> with open("tests/inputs/example_password.docx", "rb") as f:
+ ... officefile = OfficeFile(f)
+ ... officefile.load_key(password="Password1234_", verify_password=True)
+
+ >>> with open("README.md", "rb") as f:
+ ... officefile = OfficeFile(f)
+ Traceback (most recent call last):
+ ...
+ msoffcrypto.exceptions.FileFormatError: ...
+
+ >>> with open("tests/inputs/example_password.docx", "rb") as f:
+ ... officefile = OfficeFile(f)
+ ... officefile.load_key(password="0000", verify_password=True)
+ Traceback (most recent call last):
+ ...
+ msoffcrypto.exceptions.InvalidKeyError: ...
+
  Given file handle will not be closed, the file position will most certainly
  change.
  """
@@ -31,7 +50,7 @@ def OfficeFile(file):
 
  return OOXMLFile(file)
  else:
- raise Exception("Unsupported file format")
+ raise exceptions.FileFormatError("Unsupported file format")
 
  # TODO: Make format specifiable by option in case of obstruction
  # Try this first; see https://github.com/nolze/msoffcrypto-tool/issues/17
@@ -58,4 +77,4 @@ def OfficeFile(file):
 
  return Ppt97File(file)
  else:
- raise Exception("Unrecognized file format")
+ raise exceptions.FileFormatError("Unrecognized file format")

msoffcrypto/__main__.py

@@ -7,6 +7,7 @@
 
 from . import __version__
 from . import OfficeFile
+from . import exceptions
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
@@ -62,7 +63,7 @@ def main():
  return
 
  if not olefile.isOleFile(args.infile):
- raise Exception("Not OLE file")
+ raise exceptions.FileFormatError("Not OLE file")
 
  file = OfficeFile(args.infile)
 

msoffcrypto/exceptions.py

@@ -0,0 +1,22 @@
+class FileFormatError(Exception):
+ """Raised when the format of given file is unsupported or unrecognized.
+ """
+ pass
+
+
+class ParseError(Exception):
+ """Raised when the file cannot be parsed correctly.
+ """
+ pass
+
+
+class DecryptionError(Exception):
+ """Raised when the file cannot be decrypted.
+ """
+ pass
+
+
+class InvalidKeyError(DecryptionError):
+ """Raised when the given password or key is incorrect or cannot be verified.
+ """
+ pass

msoffcrypto/format/doc97.py

@@ -4,6 +4,7 @@
 
 import olefile
 
+from .. import exceptions
 from . import base
 from .common import _parse_encryptionheader, _parse_encryptionverifier
 from ..method.rc4 import DocumentRC4
@@ -274,6 +275,21 @@ def _parse_header_RC4CryptoAPI(encryptionHeader):
 
 
 class Doc97File(base.BaseOfficeFile):
+ """Return a MS-DOC file object.
+
+ Examples:
+ >>> with open("tests/inputs/rc4cryptoapi_password.doc", "rb") as f:
+ ... officefile = Doc97File(f)
+ ... officefile.load_key(password="Password1234_")
+
+ >>> with open("tests/inputs/rc4cryptoapi_password.doc", "rb") as f:
+ ... officefile = Doc97File(f)
+ ... officefile.load_key(password="0000")
+ Traceback (most recent call last):
+ ...
+ msoffcrypto.exceptions.InvalidKeyError: ...
+ """
+
  def __init__(self, file):
  self.file = file
  ole = olefile.OleFileIO(file) # do not close this, would close file
@@ -319,7 +335,7 @@ def load_key(self, password=None):
  self.key = password
  self.salt = info["salt"]
  else:
- raise Exception("Failed to verify password")
+ raise exceptions.InvalidKeyError("Failed to verify password")
  elif vMajor in [0x0002, 0x0003, 0x0004] and vMinor == 0x0002: # RC4 CryptoAPI
  info = _parse_header_RC4CryptoAPI(encryptionHeader)
  if DocumentRC4CryptoAPI.verifypw(
@@ -330,9 +346,9 @@ def load_key(self, password=None):
  self.salt = info["salt"]
  self.keySize = info["keySize"]
  else:
- raise Exception("Failed to verify password")
+ raise exceptions.InvalidKeyError("Failed to verify password")
  else:
- raise Exception("Unsupported encryption method")
+ raise exceptions.DecryptionError("Unsupported encryption method")
 
  def decrypt(self, ofile):
  # fd, _ofile_path = tempfile.mkstemp()

msoffcrypto/format/ooxml.py

@@ -6,6 +6,7 @@
 
 import olefile
 
+from .. import exceptions
 from . import base
 from .common import _parse_encryptionheader, _parse_encryptionverifier
 from ..method.ecma376_agile import ECMA376Agile
@@ -76,10 +77,25 @@ def _parseinfo(ole):
  elif versionMajor in [2, 3, 4] and versionMinor == 2: # Standard
  return "standard", _parseinfo_standard(ole)
  elif versionMajor in [3, 4] and versionMinor == 3: # Extensible
- raise Exception("Unsupported EncryptionInfo version (Extensible Encryption)")
+ raise exceptions.DecryptionError("Unsupported EncryptionInfo version (Extensible Encryption)")
 
 
 class OOXMLFile(base.BaseOfficeFile):
+ """Return an OOXML file object.
+
+ Examples:
+ >>> with open("tests/inputs/example_password.docx", "rb") as f:
+ ... officefile = OOXMLFile(f)
+ ... officefile.load_key(password="Password1234_", verify_password=True)
+
+ >>> with open("tests/inputs/example_password.docx", "rb") as f:
+ ... officefile = OOXMLFile(f)
+ ... officefile.load_key(password="0000", verify_password=True)
+ Traceback (most recent call last):
+ ...
+ msoffcrypto.exceptions.InvalidKeyError: ...
+ """
+
  def __init__(self, file):
  self.format = "ooxml"
  file.seek(0) # TODO: Investigate the effect (required for olefile.isOleFile)
@@ -104,7 +120,7 @@ def __init__(self, file):
  self.type, self.info = None, None
  self.secret_key = None
  else:
- raise Exception("Unsupported file format")
+ raise exceptions.FileFormatError("Unsupported file format")
 
  def load_key(self, password=None, private_key=None, secret_key=None, verify_password=False):
  if password:
@@ -128,7 +144,7 @@ def load_key(self, password=None, private_key=None, secret_key=None, verify_pass
  self.info["passwordKeyBits"],
  )
  if not verified:
- raise Exception("Key verification failed")
+ raise exceptions.InvalidKeyError("Key verification failed")
  elif self.type == "standard":
  self.secret_key = ECMA376Standard.makekey_from_password(
  password,
@@ -144,14 +160,14 @@ def load_key(self, password=None, private_key=None, secret_key=None, verify_pass
  self.secret_key, self.info["verifier"]["encryptedVerifier"], self.info["verifier"]["encryptedVerifierHash"]
  )
  if not verified:
- raise Exception("Key verification failed")
+ raise exceptions.InvalidKeyError("Key verification failed")
  elif self.type == "extensible":
  pass
  elif private_key:
  if self.type == "agile":
  self.secret_key = ECMA376Agile.makekey_from_privkey(private_key, self.info["encryptedKeyValue"])
  else:
- raise Exception("Unsupported key type for the encryption method")
+ raise exceptions.DecryptionError("Unsupported key type for the encryption method")
  elif secret_key:
  self.secret_key = secret_key
 
@@ -169,7 +185,7 @@ def decrypt(self, ofile, verify_integrity=False):
  stream,
  )
  if not verified:
- raise Exception("Payload integrity verification failed")
+ raise exceptions.InvalidKeyError("Payload integrity verification failed")
 
  obuf = ECMA376Agile.decrypt(self.secret_key, self.info["keyDataSalt"], self.info["keyDataHashAlgorithm"], stream)
  ofile.write(obuf)
@@ -180,7 +196,7 @@ def decrypt(self, ofile, verify_integrity=False):
 
  # If the file is successfully decrypted, there must be a valid OOXML file, i.e. a valid zip file
  if not zipfile.is_zipfile(io.BytesIO(obuf)):
- raise Exception("The file could not be decrypted with this password")
+ raise exceptions.InvalidKeyError("The file could not be decrypted with this password")
 
  def is_encrypted(self):
  # Heuristic

msoffcrypto/format/ppt97.py

@@ -4,6 +4,7 @@
 
 import olefile
 
+from .. import exceptions
 from . import base
 from .common import _parse_encryptionheader, _parse_encryptionverifier
 from ..method.rc4_cryptoapi import DocumentRC4CryptoAPI
@@ -518,6 +519,21 @@ def _parse_header_RC4CryptoAPI(encryptionInfo):
 
 
 class Ppt97File(base.BaseOfficeFile):
+ """Return a MS-PPT file object.
+
+ Examples:
+ >>> with open("tests/inputs/rc4cryptoapi_password.ppt", "rb") as f:
+ ... officefile = Ppt97File(f)
+ ... officefile.load_key(password="Password1234_")
+
+ >>> with open("tests/inputs/rc4cryptoapi_password.ppt", "rb") as f:
+ ... officefile = Ppt97File(f)
+ ... officefile.load_key(password="0000")
+ Traceback (most recent call last):
+ ...
+ msoffcrypto.exceptions.InvalidKeyError: ...
+ """
+
  def __init__(self, file):
  self.file = file
  ole = olefile.OleFileIO(file) # do not close this, would close file
@@ -580,7 +596,7 @@ def load_key(self, password=None):
  self.salt = info["salt"]
  self.keySize = info["keySize"]
  else:
- raise Exception("Failed to verify password")
+ raise exceptions.InvalidKeyError("Failed to verify password")
 
  def decrypt(self, ofile):
  # Current User Stream

msoffcrypto/format/xls97.py

@@ -4,6 +4,7 @@
 
 import olefile
 
+from .. import exceptions
 from . import base
 from .common import _parse_encryptionheader, _parse_encryptionverifier
 from ..method.rc4 import DocumentRC4
@@ -425,7 +426,7 @@ def skip_to(self, target):
  while True:
  h = self.data.read(4)
  if not h:
- raise Exception("Record not found")
+ raise exceptions.ParseError("Record not found")
  num, size = unpack("<HH", h)
  if num == target:
  return num, size
@@ -443,6 +444,21 @@ def iter_record(self):
 
 
 class Xls97File(base.BaseOfficeFile):
+ """Return a MS-XLS file object.
+
+ Examples:
+ >>> with open("tests/inputs/rc4cryptoapi_password.xls", "rb") as f:
+ ... officefile = Xls97File(f)
+ ... officefile.load_key(password="Password1234_")
+
+ >>> with open("tests/inputs/rc4cryptoapi_password.xls", "rb") as f:
+ ... officefile = Xls97File(f)
+ ... officefile.load_key(password="0000")
+ Traceback (most recent call last):
+ ...
+ msoffcrypto.exceptions.InvalidKeyError: ...
+ """
+
  def __init__(self, file):
  self.file = file
  ole = olefile.OleFileIO(file) # do not close this, would close file
@@ -488,7 +504,7 @@ def load_key(self, password=None):
  (wEncryptionType,) = unpack("<H", workbook.data.read(2))
 
  if wEncryptionType == 0x0000: # XOR obfuscation
- raise Exception("Unsupported encryption method")
+ raise exceptions.DecryptionError("Unsupported encryption method")
  elif wEncryptionType == 0x0001: # RC4
  pass
 
@@ -504,7 +520,7 @@ def load_key(self, password=None):
  self.key = password
  self.salt = info["salt"]
  else:
- raise Exception("Failed to verify password")
+ raise exceptions.InvalidKeyError("Failed to verify password")
  elif vMajor in [0x0002, 0x0003, 0x0004] and vMinor == 0x0002: # RC4 CryptoAPI
  info = _parse_header_RC4CryptoAPI(encryptionInfo)
  if DocumentRC4CryptoAPI.verifypw(
@@ -515,9 +531,9 @@ def load_key(self, password=None):
  self.salt = info["salt"]
  self.keySize = info["keySize"]
  else:
- raise Exception("Failed to verify password")
+ raise exceptions.InvalidKeyError("Failed to verify password")
  else:
- raise Exception("Unsupported encryption method")
+ raise exceptions.DecryptionError("Unsupported encryption method")
 
  def decrypt(self, ofile):
  # fd, _ofile_path = tempfile.mkstemp()
@@ -633,6 +649,6 @@ def is_encrypted(self):
  return True
  elif wEncryptionType == 0x0000: # XOR obfuscation
  # If not compatible no point stating that
- raise Exception("Unsupported encryption method")
+ raise exceptions.DecryptionError("Unsupported encryption method")
  else:
  return False