Some exceptions handled and some comments added to the files.

InvisibleCharm/__init__.py

@@ -8,7 +8,9 @@
 from log21 import get_colors as _gc
 from InvisibleCharm.lib.Console import exit as _exit
 
-__version__ = "2.2.0"
+__version__ = "2.3.0"
+__author__ = "CodeWriter21 (Mehrad Pooryoussof)"
+__github__ = "Https://GitHub.com/MPCodeWriter21/InvisibleCharm"
 
 
 def entry_point():

InvisibleCharm/__main__.py

@@ -9,6 +9,9 @@
 from InvisibleCharm.lib.Console import logger, input, verbose, quiet, exit
 from InvisibleCharm.lib.operations import win_embed, win_extract, win_attrib_hide, win_attrib_reveal, to_image_file, \
     from_image_file, embed_file, extract_file
+from InvisibleCharm.lib.Exceptions import CoverDataTypeNotFoundError, InvalidCoverDataTypeError, \
+    NoEmbeddedFileFoundError
+from InvisibleCharm.lib.File import get_names
 
 
 # Main function of script
@@ -107,12 +110,26 @@ def main():
             if args.cover:
                 logger.warn(gc("lr") + ' ! Warning: ' + gc("blm", "gr") +
                             "`to image` operation doesn't use cover file." + gc("rst"))
-            to_image_file(args.source, args.destination, args.delete, args.compress, args.encryption_pass, args.image_mode)
+            to_image_file(args.source, args.destination, args.delete, args.compress, args.encryption_pass,
+                          args.image_mode)
         elif args.embed:
             if not args.cover or not os.path.exists(args.cover) or not os.path.isfile(args.cover):
                 exit(gc("lr") + ' ! Error: Embed operation needs a cover file' +
                      '\n + Source file must be an existing file!')
-            embed_file(args.source, args.cover, args.destination, args.delete, args.compress, args.encryption_pass)
+            try:
+                embed_file(args.source, args.cover, args.destination, args.delete, args.compress, args.encryption_pass)
+            except CoverDataTypeNotFoundError:
+                confirm = input(gc("lr") + f" ! Error: Couldn't identify cover file type!\n" + gc("ly") +
+                                ' * Do you still want to use this cover file?' + gc("lw") + '(' + gc("ly") +
+                                f'Enter {gc("lm")}Y{gc("ly")} to confirm{gc("lw")}){gc("lr")}: '
+                                + gc("lg")).lower()
+                if confirm == 'y':
+                    embed_file(args.source, args.cover, args.destination, args.delete, args.compress,
+                               args.encryption_pass, True)
+                else:
+                    exit()
+            except InvalidCoverDataTypeError as ex:
+                exit('\r' + gc("lr") + " ! Error:", str(ex))
     elif args.mode.lower() in ['reveal', 'r']:
         # Checks the chosen operation and calls the suitable function
         if args.cover:
@@ -126,7 +143,22 @@ def main():
                         gc("ly") + f'Are you sure you want to remove source file({args.source})?(y/N) ' +
                         gc("lg"))
                 args.delete = answer == 'y'
-            win_extract(args.source, args.destination, args.delete, args.compress, args.encryption_pass)
+            # Gets the list of available embedded names in destination path
+            possible_names = get_names(args.source)
+            name = ''
+            if len(possible_names) == 0:
+                # Exits if no embedded file is available
+                exit(gc("lr") + f' ! Error: No win-embedded file found!')
+            elif len(possible_names) == 1:
+                # Choose the only embedded file automatically
+                name = possible_names[0]
+            else:
+                # Asks user to choose one of the embedded files
+                while name not in possible_names:
+                    logger.print(gc("lb") + 'Available names' + gc("lr") + ': ' + gc("lg") +
+                                 f'{gc("lr")}, {gc("lg")}'.join(possible_names))
+                    name = input(gc("ly") + f'Enter the name of embedded file' + gc("lr") + ': ' + gc("lg"))
+            win_extract(args.source, args.destination, args.delete, args.compress, name, args.encryption_pass)
         elif args.win_attrib:
             if args.destination:
                 logger.warn(gc("lr") + ' ! Warning: ' + gc("blm", "gr") +
@@ -135,7 +167,10 @@ def main():
         elif args.to_image:
             from_image_file(args.source, args.destination, args.delete, args.compress, args.encryption_pass)
         elif args.embed:
-            extract_file(args.source, args.destination, args.delete, args.compress, args.encryption_pass)
+            try:
+                extract_file(args.source, args.destination, args.delete, args.compress, args.encryption_pass)
+            except NoEmbeddedFileFoundError:
+                exit(gc("lr") + f' ! Error: No embedded file found!')
     else:
         exit(gc("lr") + f' ! Error: Mode: `{gc("lw")}{args.mode}{gc("lr")}` not found!')
 

InvisibleCharm/lib/Console.py

@@ -12,12 +12,25 @@
 
 
 def input(*args, end='') -> str:
+    """
+    Prints the input arguments and returns the user input.
+
+    :param args: Input arguments to write in the console.
+    :param end:
+    :return: str
+    """
     logger.info(*args, end=end)
     return _input_backup('')
 
 
 # Exits
 def exit(*args) -> None:
+    """
+    Prints the input arguments and exits the program.
+
+    :param args: Input arguments to write in the console.
+    :return: None
+    """
     if args:
         logger.error(*args)
     logger.error(end='\033[0m')
@@ -26,11 +39,21 @@ def exit(*args) -> None:
 
 # Enables verbose mode
 def verbose() -> int:
+    """
+    Enables verbose mode.
+
+    :return: int: logger.level
+    """
     logger.setLevel(_DEBUG)
     return logger.level
 
 
 # Enables quiet mode
 def quiet() -> int:
+    """
+    Enables quiet mode.
+
+    :return: int: logger.level
+    """
     logger.setLevel(_WARNING)
     return logger.level

InvisibleCharm/lib/Exceptions.py

@@ -0,0 +1,37 @@
+# InvisibleCharm.lib.Exceptions.py
+# CodeWriter21
+
+from zlib import error as _zlib_error
+from magic.magic import MagicException as _MagicException
+
+__all__ = ['CouldNotDecompressError', 'InvalidCoverDataTypeError', 'CoverDataTypeNotFoundError',
+           'NoEmbeddedDataFoundError', 'WinEmbeddedFileNotFoundError', 'NoEmbeddedFileFoundError',
+           'NoWinEmbeddedFileFoundError']
+
+
+class CouldNotDecompressError(_zlib_error):
+    pass
+
+
+class InvalidCoverDataTypeError(_MagicException):
+    pass
+
+
+class CoverDataTypeNotFoundError(InvalidCoverDataTypeError):
+    pass
+
+
+class NoEmbeddedDataFoundError:
+    pass
+
+
+class NoEmbeddedFileFoundError(NoEmbeddedDataFoundError, FileNotFoundError):
+    pass
+
+
+class WinEmbeddedFileNotFoundError(FileNotFoundError):
+    pass
+
+
+class NoWinEmbeddedFileFoundError(WinEmbeddedFileNotFoundError, NoEmbeddedFileFoundError):
+    pass

InvisibleCharm/lib/File.py

@@ -2,15 +2,38 @@
 # CodeWriter21
 
 import os as _os
+from log21 import get_colors as _gc
+from typing import Union as _Union, List as _List
 # We use getoutput function to get the output of a command
 from subprocess import getoutput as _getoutput
-from log21 import get_colors as _gc
 from InvisibleCharm.lib.Console import logger as _logger
 from InvisibleCharm.lib.data.Prepare import prepare_data as _prepare_data, add_num as _add_num
 
 
 # Opens a file and returns prepared content
-def open_file(path: str, name: str, hiding: bool = True, compress: bool = False, encrypt_pass: str = None) -> bytes:
+def open_file(path: str, name: str, hiding: bool = True, compress: bool = False,
+              encrypt_pass: _Union[str, bytes] = '') -> bytes:
+    """
+    Gets a file path and a name and reads the contents of the file and prepares it using
+    InvisibleCharm.lib.data.Prepare.prepare_data function.
+
+    :param path: str: Input file path
+    :param name: str: A name to use for debug messages
+    :param hiding: bool = True: Are you going to hide this file?
+    :param compress: bool = False: Do you want to compress this file?
+    :param encrypt_pass: Union[str, bytes] = '': A password for encrypting the file
+    :return: bytes: Prepared data
+    """
+
+    # Checks whether the inputs are valid
+    if not isinstance(path, str):
+        raise TypeError('`path` must be an instance of str.')
+    if not _os.path.exists(path):
+        raise FileNotFoundError('`path` must be an existing file path.')
+    if not isinstance(encrypt_pass, (str, bytes)) and encrypt_pass:
+        raise TypeError('`encrypt_pass` must be an instance of str or bytes.')
+
+    name = str(name)
     _logger.debug(_gc("ly") + f' * Reading {name.lower()} file...', end='')
     with open(path, 'rb') as file:
         data = file.read()
@@ -19,7 +42,23 @@ def open_file(path: str, name: str, hiding: bool = True, compress: bool = False,
 
 
 # Writes data in the given path
-def save_file(path: str, data: bytes) -> None:
+def save_file(path: str, data: _Union[bytes, str]) -> None:
+    """
+    Writes data in the given path.
+
+    :param path: str: File path to write the data.
+    :param data: Union[bytes, str]: Content to write.
+    :return: None
+    """
+
+    # Checks whether the inputs are valid
+    if not isinstance(path, str):
+        raise TypeError('`path` must be an instance of str.')
+    if not isinstance(data, (str, bytes)):
+        raise TypeError('`encrypt_pass` must be an instance of str or bytes.')
+    if isinstance(data, str):
+        data = data.encode()
+
     _logger.debug(_gc("ly") + ' * Writing in destination file...', end='')
     try:
         with open(path, 'wb') as dest_file:
@@ -37,13 +76,39 @@ def save_file(path: str, data: bytes) -> None:
 
 # Deletes the file in the given path
 def delete_source_file(path: str) -> None:
+    """
+    Deletes the file in the given path.
+
+    :param path: str: File path to write the data.
+    :return: None
+    """
+
+    # Checks whether the inputs are valid
+    if not isinstance(path, str):
+        raise TypeError('`path` must be an instance of str.')
+    if not _os.path.exists(path):
+        raise FileNotFoundError('`path` must be an existing file path.')
+
     _logger.debug(_gc("ly") + ' * Deleting source file...', end='')
     _os.remove(path)
     _logger.debug('\r' + _gc("ly") + ' = Source file deleted')
 
 
-# Returns the possible names of the embedded files in a windows path
-def get_names(path: str) -> list:
+# Returns the possible names of the embedded files in a Windows path
+def get_names(path: str) -> _List[str]:
+    """
+    Returns the possible names of the embedded files in a Windows path.
+
+    :param path: str: File path to write the data.
+    :return: List[str]
+    """
+
+    # Checks whether the inputs are valid
+    if not isinstance(path, str):
+        raise TypeError('`path` must be an instance of str.')
+    if not _os.path.exists(path):
+        raise FileNotFoundError('`path` must be an existing file path.')
+
     output = _getoutput('dir ' + _os.path.split(path)[0] + ' /r /a')
     names = []
     for line in output.split('\n'):

InvisibleCharm/lib/data/Encryption.py

@@ -3,19 +3,52 @@
 
 # We use hashlib library to hash passwords
 import hashlib
+from log21 import get_colors as _gc
+from typing import Union as _Union, Tuple as _Tuple
 # We use _AES to encrypt and decrypt data
 from Crypto.Cipher import AES as _AES
-from log21 import get_colors as _gc
 from InvisibleCharm.lib.Console import logger as _logger
 
 __all__ = ['encrypt', 'decrypt']
 
 
-# Encrypts data using AES and costume password
-def encrypt(data: bytes, password: str) -> bytes:
+def __prepare(data: _Union[str, bytes], password: _Union[str, bytes]) -> _Tuple[bytes, bytes, bytes]:
+    """
+    Prepares data for encrypt and decrypt functions.
+
+    :param data: Union[str, bytes]
+    :param password: Union[str, bytes]
+    :return: Tuple[bytes, bytes, bytes]: data, md5_password_hash, sha512_password_hash
+    """
+    # Checks whether the inputs are valid
+    if not isinstance(data, (str, bytes)):
+        raise TypeError('`data` must be an instance of str or bytes.')
+    if isinstance(data, str):
+        data = data.encode()
+    if not isinstance(password, (str, bytes)):
+        raise TypeError('`password` must be an instance of str or bytes.')
+    if isinstance(password, str):
+        password = password.encode()
+
     _logger.debug(_gc("ly") + ' * Hashing password...')
-    md5 = hashlib.md5(password.encode()).digest()
-    sha512 = hashlib.sha512(password.encode()).digest()
+    md5 = hashlib.md5(password).digest()
+    sha512 = hashlib.sha512(password).digest()
+
+    return data, md5, sha512
+
+
+# Encrypts data using AES and costume password
+def encrypt(data: _Union[str, bytes], password: _Union[str, bytes]) -> bytes:
+    """
+    Encrypts data using AES and costume password.
+
+    :param data: Union[str, bytes]: Data to encrypt.
+    :param password: Password for encryption.
+    :return: bytes: Encrypted data.
+    """
+
+    data, md5, sha512 = __prepare(data, password)
+
     _logger.debug(_gc("ly") + ' * Encrypting data...', end='')
     cipher = _AES.new(md5, _AES.MODE_EAX, nonce=sha512)
     data = cipher.encrypt(data)
@@ -24,10 +57,17 @@ def encrypt(data: bytes, password: str) -> bytes:
 
 
 # Decrypts data using AES and costume password
-def decrypt(data: bytes, password: str) -> bytes:
-    _logger.debug(_gc("ly") + ' * Hashing password...')
-    md5 = hashlib.md5(password.encode()).digest()
-    sha512 = hashlib.sha512(password.encode()).digest()
+def decrypt(data: _Union[str, bytes], password: _Union[str, bytes]) -> bytes:
+    """
+    Decrypts data using AES and costume password.
+
+    :param data: Union[str, bytes]: Data to decrypt.
+    :param password: Password for decryption.
+    :return: bytes: Decrypted data.
+    """
+
+    data, md5, sha512 = __prepare(data, password)
+
     _logger.debug(_gc("ly") + ' * Decrypting data...', end='')
     cipher = _AES.new(md5, _AES.MODE_EAX, nonce=sha512)
     data = cipher.decrypt(data)