From 4e5dc7485e808aaf981c0107a5c93bd8271b6744 Mon Sep 17 00:00:00 2001
From: Ivan Maslov <Ivan.Maslov@UnicodeLabs.ru>
Date: Sun, 1 Mar 2020 20:46:43 +0300
Subject: [PATCH] #SafeSource Create tested # Run in prototype

---
 Sources/pyOpenRPA/Tools/SafeSource/Crypter.py | 22 ++++++-
 .../pyOpenRPA/Tools/SafeSource/DistrCreate.py | 58 ++++++++++++++-----
 .../pyOpenRPA/Tools/SafeSource/DistrRun.py    | 44 +++++++-------
 .../pyOpenRPA/Tools/SafeSource/__main__.py    |  2 +-
 4 files changed, 87 insertions(+), 39 deletions(-)

diff --git a/Sources/pyOpenRPA/Tools/SafeSource/Crypter.py b/Sources/pyOpenRPA/Tools/SafeSource/Crypter.py
index 62705ef2..6a44c717 100644
--- a/Sources/pyOpenRPA/Tools/SafeSource/Crypter.py
+++ b/Sources/pyOpenRPA/Tools/SafeSource/Crypter.py
@@ -34,7 +34,6 @@ def encrypt_file(key, in_filename, out_filename=None, chunksize=64*1024):
             outfile.write(iv)
             while True:
                 chunk = infile.read(chunksize)
-                print(chunk)
                 if len(chunk) == 0:
                     break
                 elif len(chunk) % 16 != 0:
@@ -61,4 +60,23 @@ def decrypt_file(key, in_filename, out_filename=None, chunksize=24*1024):
                 if len(chunk) == 0:
                     break
                 outfile.write(decryptor.decrypt(chunk))
-            outfile.truncate(origsize)
\ No newline at end of file
+            outfile.truncate(origsize)
+def decrypt_file_bytes(key, in_filename, chunksize=24*1024):
+    """ Расшифровывает файл, используя AES (режим CBC) с
+        заданным ключом.
+    """
+    lResult = b''
+    if not out_filename:
+        out_filename = os.path.splitext(in_filename)[0]
+    with open(in_filename, 'rb') as infile:
+        origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
+        iv = infile.read(16)
+        decryptor = AES.new(key, AES.MODE_CBC, iv)
+        with open(out_filename, 'wb') as outfile:
+            while True:
+                chunk = infile.read(chunksize)
+                if len(chunk) == 0:
+                    break
+                lResult = lResult+decryptor.decrypt(chunk)
+            lResult=lResult[0:origsize-1]
+    return lResult
\ No newline at end of file
diff --git a/Sources/pyOpenRPA/Tools/SafeSource/DistrCreate.py b/Sources/pyOpenRPA/Tools/SafeSource/DistrCreate.py
index 277d6563..d8d4f322 100644
--- a/Sources/pyOpenRPA/Tools/SafeSource/DistrCreate.py
+++ b/Sources/pyOpenRPA/Tools/SafeSource/DistrCreate.py
@@ -3,26 +3,52 @@ import hashlib # Create hash of the key
 import tkinter #tkinter
 from tkinter import filedialog
 from . import Crypter
-
+import shutil # Copy folder
+import os # rm dir, listdir
+import glob # list files
+import datetime #Get current datetime
+from Crypto.Cipher import AES
+# Settings
+gInUncryptedExtension = "py" # cry for filename.cry
+gOutCryptedExtension = "cry" # cry for filename.cry
+gFileMaskToDelete = "pyc" #Remove all .pyc files
 # Create process run
 def Run():
-    # Step 1 - select folder to crypt
+    print(f"{str(datetime.datetime.now())}: Start to create crypted distr")
+    ############# Step 1 - select folder to copy 
     lStep_1_TkinterRoot = tkinter.Tk()
     lStep_1_TkinterRoot.withdraw()
-    #lStep_1_FolderPath = filedialog.askdirectory(parent=lStep_1_TkinterRoot,title='Please select a source code directory to crypt')
-    
-    # Step 2 - set password to crypt
+    lStep_1_FolderPath = filedialog.askdirectory(parent=lStep_1_TkinterRoot,title='Please select a source code directory to crypt')
+    ############# Step 2 - Select folder to save
+    lStep_2_FolderPath = filedialog.askdirectory(parent=lStep_1_TkinterRoot,title='Please select a folder to save the crypted result')	
+    # Before delete the folder check if directory is empty - else ask user to confirm the deletion
+    if len(os.listdir(lStep_2_FolderPath)) != 0:
+        lStep_2_ConfirmDelete = pyautogui.confirm('The destination folder contains some files/folders. The are will be removed. Continue?', buttons=['Yes', 'No'])
+        if lStep_2_ConfirmDelete=="No":
+            raise Exception("Stop program - user suggestion. Don't want to clear destination folder")
+    shutil.rmtree(lStep_2_FolderPath, ignore_errors=False, onerror=None)
+    ############# Step 3 - Copy folder
+    shutil.copytree(lStep_1_FolderPath, lStep_2_FolderPath)
+    ############# Step 3.1 - Remove files to delete 
+    lPyFilesToDeleteList = [f for f in glob.glob(os.path.join(lStep_2_FolderPath,f"**/*.{gFileMaskToDelete}"), recursive=True)]
+    for lFileItem in lPyFilesToDeleteList:
+        #Create right \\ splashes
+        lFileItem = os.path.abspath(lFileItem)
+        #Remove old file
+        os.remove(lFileItem)
+    ############# Step 4 - Get file list with extension .py
+    lPyFileList = [f for f in glob.glob(os.path.join(lStep_2_FolderPath,f"**/*.{gInUncryptedExtension}"), recursive=True)]
+    ############# Step 5 - Ask and confirm the secret word
     lKeyHashStr_1 = hashlib.sha256(pyautogui.password('Please enter the key to protect source code').encode("utf-8")).digest()
-    print(lKeyHashStr_1)
     lKeyHashStr_2 = hashlib.sha256(pyautogui.password('Please repeat the key to protect source code').encode("utf-8")).digest()
-    print(lKeyHashStr_2)
     if lKeyHashStr_1 == lKeyHashStr_2:
-        print("Passed - do AES encryption")
-        from Crypto.Cipher import AES
-        #Crypter.encrypt_file(lKeyHashStr_1, "c:\\test.png", "c:\\test.enc")
-        Crypter.decrypt_file(lKeyHashStr_1, "c:\\test.enc", "c:\\testEnc2.png")
-        #IV = 16 * '\x00'           # Вектор инициализации: обсуждается позже
-        #mode = AES.MODE_CBC
-        #encryptor = AES.new(lKeyHashStr_1, mode, IV=IV)
-        #text = "Text to be crypted!"
-        #ciphertext = encryptor.encrypt(text)
+        for lFileItem in lPyFileList:
+            #Create right \\ splashes
+            lFileItem = os.path.abspath(lFileItem)
+            Crypter.encrypt_file(lKeyHashStr_1, lFileItem, f"{lFileItem[0:-2]}{gOutCryptedExtension}")
+            #Remove old file
+            os.remove(lFileItem)
+    else:
+        raise Exception("User set different secret key 1 and key 2")
+    ############ Step 6 - Final stage
+    print(f"{str(datetime.datetime.now())}: Crypted distr is created!")
\ No newline at end of file
diff --git a/Sources/pyOpenRPA/Tools/SafeSource/DistrRun.py b/Sources/pyOpenRPA/Tools/SafeSource/DistrRun.py
index afa677d6..3e7d2b40 100644
--- a/Sources/pyOpenRPA/Tools/SafeSource/DistrRun.py
+++ b/Sources/pyOpenRPA/Tools/SafeSource/DistrRun.py
@@ -4,17 +4,15 @@ import os
 import sys
 import imp
 import base64
-
-EXT = '.b64'
+from . import Crypter # Crypto functions
+EXT = '.cry'
 
 # How to run
 # sys.meta_path.append(Base64Importer(root_pkg_path))
 
 # Init cryptographer
 def CryptographerInit(inFolderPath):
-    print(1)    
     sys.meta_path.append(Base64Importer(inFolderPath))
-
 #===============================================================================
 class Base64Importer(object):
     """Служит для поиска и импорта python-модулей, кодированных в base64
@@ -22,10 +20,9 @@ class Base64Importer(object):
     Класс реализует Import Protocol (PEP 302) для возможности импортирования
     модулей, зашифрованных в base64 из указанного пакета.
     """
-
     #---------------------------------------------------------------------------
-    def __init__(self, root_package_path):
-
+    def __init__(self, root_package_path, inKey):
+        self.mKeyStr = inKey
         self.__modules_info = self.__collect_modules_info(root_package_path)
 
     #---------------------------------------------------------------------------
@@ -43,41 +40,31 @@ class Base64Importer(object):
     #---------------------------------------------------------------------------
     def load_module(self, fullname):
         """Метод загружает base64 модуль
-
         Если модуль с именем fullname является base64, то метод попытается его
         загрузить. Возбуждает исключение ImportError в случае любой ошибки.
         """
         if not fullname in self.__modules_info:
             raise ImportError(fullname)
-
         # Для потокобезопасности
         imp.acquire_lock()
-
         try:
             mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
-
             mod.__file__ = "<{}>".format(self.__class__.__name__)
             mod.__loader__ = self
-
             if self.is_package(fullname):
                 mod.__path__ = []
                 mod.__package__ = fullname
             else:
                 mod.__package__ = fullname.rpartition('.')[0]
-
             src = self.get_source(fullname)
-
             try:
                 exec(src) in mod.__dict__
             except:
                 del sys.modules[fullname]
                 raise ImportError(fullname)
-
         finally:
             imp.release_lock()
-
         return mod
-
     #---------------------------------------------------------------------------
     def is_package(self, fullname):
         """Возвращает True если fullname является пакетом
@@ -93,8 +80,9 @@ class Base64Importer(object):
         filename = self.__modules_info[fullname]['filename']
 
         try:
-            with open(filename, 'r') as ifile:
-                src = base64.decodestring(ifile.read())
+            src = Crypter.decrypt_file_bytes(key = self.mKeyStr, in_filename = filename).decode("utf-8")
+            #with open(filename, 'r') as ifile:
+            #    src = base64.decodestring(ifile.read())
         except IOError:
             src = ''
 
@@ -132,4 +120,20 @@ class Base64Importer(object):
                     'ispackage': False
                 }
 
-        return modules
\ No newline at end of file
+        return modules
+# Settings
+gInUncryptedExtension = "py" # cry for filename.cry
+gOutCryptedExtension = "cry" # cry for filename.cry
+gFileMaskToDelete = "pyc" #Remove all .pyc files
+# Create process run
+def Run():
+    print(f"{str(datetime.datetime.now())}: Run decryptography")
+    ############# Step 5 - Ask and confirm the secret word
+    lKeyHashStr_1 = hashlib.sha256(pyautogui.password('Please enter the key to protect source code').encode("utf-8")).digest()
+    lKeyHashStr_2 = hashlib.sha256(pyautogui.password('Please repeat the key to protect source code').encode("utf-8")).digest()
+    if lKeyHashStr_1 == lKeyHashStr_2:
+        CryptographerInit("", inKey = lKeyHashStr_1)
+    else:
+        raise Exception("User set different secret key 1 and key 2")
+    ############ Step 6 - Final stage
+    print(f"{str(datetime.datetime.now())}: Cryprography module has been successfully initialized")
\ No newline at end of file
diff --git a/Sources/pyOpenRPA/Tools/SafeSource/__main__.py b/Sources/pyOpenRPA/Tools/SafeSource/__main__.py
index dce42f0d..6ebdd71a 100644
--- a/Sources/pyOpenRPA/Tools/SafeSource/__main__.py
+++ b/Sources/pyOpenRPA/Tools/SafeSource/__main__.py
@@ -14,7 +14,7 @@ from pyOpenRPA.Tools.SafeSource import DistrRun
 from pyOpenRPA.Tools.SafeSource import DistrCreate
 #Mode RUN
 if sys.argv[1].upper() == "RUN":
-    DistrRun.CryptographerInit("")
+    DistrRun.Run()
 #Mode CREATE
 if sys.argv[1].upper() == "CREATE":
     DistrCreate.Run()