sFE.py

import os
import sys
import base64
import secrets
import hashlib
import logging
import platform
import ctypes
import lzma
import json
import shutil
from datetime import datetime
from typing import Tuple, Optional, Dict, Union
from pathlib import Path
from cryptography.hazmat.primitives.ciphers.aead import AESGCM, ChaCha20Poly1305
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.exceptions import InvalidTag
import argparse
import argon2

def clear():
    system = platform.system().lower()

    if system == 'windows':
        _ = os.system('cls')
    elif system == 'linux' or system == 'darwin':
        _ = os.system('clear')
    elif system == 'android':
        System.Clear()
        print("Cannot use this tool on Android Shell emulation")
        exit()

class SecurityError(Exception):
    pass

class SecureStorage:
    def __init__(self, base_path: str = None):
        self.base_path = base_path or os.path.join(os.path.expanduser("~"), ".secure_")
        self.keys_path = os.path.join(self.base_path, "keys")
        self._setup_secure_directory()
        
    def _setup_secure_directory(self):
        """Create directories with appropriate permissions."""
        os.makedirs(self.keys_path, exist_ok=True)
        
        if platform.system() == 'Windows':
            self._set_windows_permissions(self.keys_path)
        else:
            self._set_unix_permissions(self.keys_path)
            
    def _set_windows_permissions(self, path: str):
        """Set secure Windows ACLs."""
        try:
            import win32security
            import win32api
            import win32con
            
            token = win32security.OpenProcessToken(win32api.GetCurrentProcess(), win32con.TOKEN_QUERY)
            user_sid = win32security.GetTokenInformation(token, win32security.TokenUser)[0]
            
            dacl = win32security.ACL()
            dacl.AddAccessAllowedAce(
                win32security.ACL_REVISION,
                win32con.GENERIC_READ | win32con.GENERIC_WRITE,
                user_sid
            )
            
            sd = win32security.SECURITY_DESCRIPTOR()
            sd.SetSecurityDescriptorDacl(1, dacl, 0)
            win32security.SetFileSecurity(
                path,
                win32security.DACL_SECURITY_INFORMATION,
                sd
            )
        except ImportError:
            os.chmod(path, 0o700)
        
    def _set_unix_permissions(self, path: str):
        os.chmod(path, 0o600)

    def _derive_key(self, password: str, salt: bytes, purpose: str) -> bytes:
        purpose_salt = salt + purpose.encode()
        
        ph = argon2.PasswordHasher(
            time_cost=4,          
            memory_cost=1048576,  
            parallelism=8,        
            hash_len=32,          
            type=argon2.Type.ID  
        )
        
        hash_value = argon2.low_level.hash_secret_raw(
            secret=password.encode(),
            salt=purpose_salt,
            time_cost=ph.time_cost,
            memory_cost=ph.memory_cost,
            parallelism=ph.parallelism,
            hash_len=ph.hash_len,
            type=ph.type
        )
        
        return hash_value

    def store_key(self, key_data: bytes, key_id: str, password: str) -> None:
        salt = secrets.token_bytes(32)
        
        key = self._derive_key(password, salt, f"key_storage_{key_id}")
        
        chacha = ChaCha20Poly1305(key)
        nonce = secrets.token_bytes(12)
        encrypted_data = chacha.encrypt(nonce, key_data, None)
        
        key_info = {
            'salt': base64.b64encode(salt).decode('utf-8'),
            'nonce': base64.b64encode(nonce).decode('utf-8'),
            'encrypted_data': base64.b64encode(encrypted_data).decode('utf-8')
        }
        
        key_path = os.path.join(self.keys_path, f"{key_id}.key")
        temp_path = key_path + '.tmp'
        
        with open(temp_path, 'w') as f:
            json.dump(key_info, f)
            f.flush()
            os.fsync(f.fileno())
            
        os.replace(temp_path, key_path)
        
        if platform.system() == 'Windows':
            self._set_windows_permissions(key_path)
        else:
            self._set_unix_permissions(key_path)

    def retrieve_key(self, key_id: str, password: str) -> bytes:
        key_path = os.path.join(self.keys_path, f"{key_id}.key")
        
        try:
            with open(key_path, 'r') as f:
                key_info = json.load(f)
                
            salt = base64.b64decode(key_info['salt'])
            nonce = base64.b64decode(key_info['nonce'])
            encrypted_data = base64.b64decode(key_info['encrypted_data'])
            
            key = self._derive_key(password, salt, f"key_storage_{key_id}")
            chacha = ChaCha20Poly1305(key)
            
            return chacha.decrypt(nonce, encrypted_data, None)
            
        except (FileNotFoundError, KeyError, InvalidTag) as e:
            raise SecurityError(f"Failed to retrieve key: {str(e)}")

class SecureEncryptor:
    def __init__(self):
        self.logger = self._setup_logging()
        self._check_security_environment()
        self.key_storage = SecureStorage()
        
    def _setup_logging(self) -> logging.Logger:
        log_dir = Path.home() / ".secure_" / "logs"
        log_dir.mkdir(parents=True, exist_ok=True)
        
        if platform.system() != 'Windows':
            log_dir.chmod(0o700)
            
        logger = logging.getLogger("SecureEncryptor")
        logger.setLevel(logging.INFO)
        
        handler = logging.FileHandler(
            log_dir / f"encryptor_{datetime.now().strftime('%Y%m%d')}.log"
        )
        handler.setFormatter(
            logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        )
        logger.addHandler(handler)
        
        return logger

    def _check_security_environment(self):
        if self._requires_elevation():
            if platform.system() == 'Windows':
                if not ctypes.windll.shell32.IsUserAnAdmin():
                    raise SecurityError("Administrative privileges required")
            else:
                if os.geteuid() != 0:
                    raise SecurityError("Root privileges required")
        
        self._secure_memory()

    def _secure_memory(self):
        if platform.system() == 'Windows':
            kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
            kernel32.VirtualLock.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
            kernel32.VirtualLock.restype = ctypes.c_bool
            
        else:
            import resource
            
            resource.setrlimit(resource.RLIMIT_MEMLOCK, (-1, -1))
            
            try:
                import ctypes.util
                libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)
                if libc.mlockall(2):  # MCL_CURRENT | MCL_FUTURE
                    raise OSError(ctypes.get_errno())
            except Exception as e:
                self.logger.warning(f"Failed to lock memory: {str(e)}")

    def _requires_elevation(self) -> bool:
        return False

    def _generate_file_id(self, file_path: str) -> str:
        timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
        random_bytes = secrets.token_bytes(8)
        path_hash = hashlib.sha256(str(file_path).encode()).hexdigest()[:8]
        return f"{timestamp}_{path_hash}_{base64.urlsafe_b64encode(random_bytes).decode()}"

    def encrypt_file(self, file_path: str, password: str) -> str:
        try:
            file_id = self._generate_file_id(file_path)
            output_path = f"{file_path}.encrypted"
            
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=4096
            )
            public_key = private_key.public_key()
            
            private_pem = private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
            self.key_storage.store_key(private_pem, file_id, password)
            
            data_key = secrets.token_bytes(32)
            
            encrypted_key = public_key.encrypt(
                data_key,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            
            chacha = ChaCha20Poly1305(data_key)
            nonce = secrets.token_bytes(12)
            
            with open(file_path, 'rb') as f:
                file_data = f.read()
                
            compressed_data = lzma.compress(file_data)
            
            encrypted_data = chacha.encrypt(nonce, compressed_data, None)
            
            with open(output_path, 'wb') as f:
                header = {
                    'version': '2.0',
                    'file_id': file_id,
                    'nonce': base64.b64encode(nonce).decode('utf-8'),
                    'encrypted_key': base64.b64encode(encrypted_key).decode('utf-8')
                }
                header_bytes = json.dumps(header).encode()
                f.write(len(header_bytes).to_bytes(4, 'big'))
                f.write(header_bytes)
                
                f.write(encrypted_data)
            
            if platform.system() == 'Windows':
                self._set_windows_permissions(output_path)
            else:
                os.chmod(output_path, 0o600)
            
            self.logger.info(f"File encrypted successfully: {output_path}")
            return output_path
            
        except Exception as e:
            self.logger.error(f"File encryption failed: {str(e)}")
            raise

    def decrypt_file(self, encrypted_path: str, password: str) -> str:
        try:
            with open(encrypted_path, 'rb') as f:
                header_size = int.from_bytes(f.read(4), 'big')
                header = json.loads(f.read(header_size))
                encrypted_data = f.read()
            
            file_id = header['file_id']
            nonce = base64.b64decode(header['nonce'])
            encrypted_key = base64.b64decode(header['encrypted_key'])
            
            private_pem = self.key_storage.retrieve_key(file_id, password)
            private_key = serialization.load_pem_private_key(
                private_pem,
                password=None
            )
            
            data_key = private_key.decrypt(
                encrypted_key,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            
            chacha = ChaCha20Poly1305(data_key)
            decrypted_data = chacha.decrypt(nonce, encrypted_data, None)
            
            decompressed_data = lzma.decompress(decrypted_data)
            
            output_path = encrypted_path.rsplit('.encrypted', 1)[0] + '_decrypted'
            with open(output_path, 'wb') as f:
                f.write(decompressed_data)
            
            if platform.system() == 'Windows':
                self._set_windows_permissions(output_path)
            else:
                os.chmod(output_path, 0o600)
            
            self.logger.info(f"File decrypted successfully: {output_path}")
            return output_path
            
        except Exception as e:
            self.logger.error(f"File decryption failed: {str(e)}")
            raise

    def encrypt_directory(self, dir_path: str, password: str) -> str:
        try:
            dir_id = self._generate_file_id(dir_path)
            encrypted_dir = f"{dir_path}_encrypted"
            os.makedirs(encrypted_dir, exist_ok=True)
            
            metadata = {
                'version': '2.0',
                'dir_id': dir_id,
                'original_name': os.path.basename(dir_path),
                'structure': {},
                'file_ids': {}
            }
            
            for root, dirs, files in os.walk(dir_path):
                rel_path = os.path.relpath(root, dir_path)
                if rel_path == '.':
                    rel_path = ''
                
                enc_dir_path = os.path.join(encrypted_dir, rel_path)
                os.makedirs(enc_dir_path, exist_ok=True)
                
                if rel_path:
                    metadata['structure'][rel_path] = {
                        'dirs': dirs,
                        'files': files
                    }
                
                for file in files:
                    file_path = os.path.join(root, file)
                    rel_file_path = os.path.join(rel_path, file)
                    
                    file_id = self._generate_file_id(rel_file_path)
                    metadata['file_ids'][rel_file_path] = file_id
                    
                    encrypted_path = os.path.join(enc_dir_path, f"{file}.encrypted")
                    self.encrypt_file(file_path, password)
            
            metadata_path = os.path.join(encrypted_dir, "directory.meta")
            self._save_encrypted_metadata(metadata, metadata_path, password)
            
            readme_path = os.path.join(encrypted_dir, "README.txt")
            with open(readme_path, 'w') as f:
                f.write("Encrypted Directory\n")
                f.write("==================\n\n")
                f.write("This directory contains encrypted files. To decrypt:\n")
                f.write("1. Use the same encryption tool version\n")
                f.write("2. Select 'Decrypt Directory' option\n")
                f.write("3. Provide the original encryption password\n")
                f.write("\nWarning: Do not modify any files in this directory\n")
            
            if platform.system() == 'Windows':
                self._set_windows_permissions(encrypted_dir)
            else:
                self._secure_unix_directory(encrypted_dir)
            
            return encrypted_dir
            
        except Exception as e:
            self.logger.error(f"Directory encryption failed: {str(e)}")
            if os.path.exists(encrypted_dir):
                shutil.rmtree(encrypted_dir)
            raise

    def decrypt_directory(self, encrypted_dir: str, password: str) -> str:
        try:
            metadata_path = os.path.join(encrypted_dir, "directory.meta")
            metadata = self._read_encrypted_metadata(metadata_path, password)
            
            output_dir = f"{os.path.dirname(encrypted_dir)}/{metadata['original_name']}_decrypted"
            os.makedirs(output_dir, exist_ok=True)
            
            for rel_path, structure in metadata['structure'].items():
                dir_path = os.path.join(output_dir, rel_path)
                os.makedirs(dir_path, exist_ok=True)
            
            for rel_file_path, file_id in metadata['file_ids'].items():
                encrypted_path = os.path.join(encrypted_dir, rel_file_path + '.encrypted')
                decrypted_path = os.path.join(output_dir, rel_file_path)
                
                os.makedirs(os.path.dirname(decrypted_path), exist_ok=True)
                self.decrypt_file(encrypted_path, password)
            
            return output_dir
            
        except Exception as e:
            self.logger.error(f"Directory decryption failed: {str(e)}")
            if os.path.exists(output_dir):
                shutil.rmtree(output_dir)
            raise

    def _save_encrypted_metadata(self, metadata: dict, output_path: str, password: str):
        metadata_key = secrets.token_bytes(32)
        metadata_bytes = json.dumps(metadata).encode()
        chacha = ChaCha20Poly1305(metadata_key)
        nonce = secrets.token_bytes(12)
        encrypted_metadata = chacha.encrypt(nonce, metadata_bytes, None)
        
        metadata_id = self._generate_file_id("metadata")
        self.key_storage.store_key(metadata_key, metadata_id, password)
        
        with open(output_path, 'wb') as f:
            header = {
                'metadata_id': metadata_id,
                'nonce': base64.b64encode(nonce).decode('utf-8')
            }
            header_bytes = json.dumps(header).encode()
            f.write(len(header_bytes).to_bytes(4, 'big'))
            f.write(header_bytes)
            f.write(encrypted_metadata)

    def _read_encrypted_metadata(self, input_path: str, password: str) -> dict:
        with open(input_path, 'rb') as f:
            header_size = int.from_bytes(f.read(4), 'big')
            header = json.loads(f.read(header_size))
            encrypted_metadata = f.read()
        
        metadata_key = self.key_storage.retrieve_key(header['metadata_id'], password)
        nonce = base64.b64decode(header['nonce'])
        
        chacha = ChaCha20Poly1305(metadata_key)
        metadata_bytes = chacha.decrypt(nonce, encrypted_metadata, None)
        return json.loads(metadata_bytes)

    def _secure_unix_directory(self, dir_path: str):
        for root, dirs, files in os.walk(dir_path):
            os.chmod(root, 0o700)  
            for file in files:
                os.chmod(os.path.join(root, file), 0o600) 

    def secure_wipe(self, path: str):
        if os.path.isfile(path):
            self._secure_wipe_file(path)
        elif os.path.isdir(path):
            self._secure_wipe_directory(path)

    def _secure_wipe_file(self, file_path: str):
        try:
            file_size = os.path.getsize(file_path)
            
            with open(file_path, 'wb') as f:
                # Multiple pass secure wiping
                patterns = [
                    b'\x00' * 1024,  # zeros
                    b'\xFF' * 1024,  # ones
                    secrets.token_bytes(1024)  # random
                ]
                
                for pattern in patterns:
                    for _ in range(0, file_size, 1024):
                        remaining = min(1024, file_size - f.tell())
                        f.write(pattern[:remaining])
                        f.flush()
                        os.fsync(f.fileno())
            
            os.remove(file_path)
            
        except Exception as e:
            self.logger.error(f"Secure file wiping failed: {str(e)}")
            raise

    def _secure_wipe_directory(self, dir_path: str):
        try:
            for root, dirs, files in os.walk(dir_path, topdown=False):
                for name in files:
                    self._secure_wipe_file(os.path.join(root, name))
                for name in dirs:
                    os.rmdir(os.path.join(root, name))
            os.rmdir(dir_path)
            
        except Exception as e:
            self.logger.error(f"Secure directory wiping failed: {str(e)}")
            raise

def main():
    clear()
    print("""
               ----------------------------------------------
                     ~T.R.U.S.T Secure File Encryption~
                Military grade encryption at your fingertips
               ----------------------------------------------
    """)
    parser = argparse.ArgumentParser(description='Help:')
    parser.add_argument('action', choices=['encrypt', 'decrypt'], help='Action to perform')
    parser.add_argument('path', help='File or directory path')
    parser.add_argument('--password', help='Encryption/decryption password')
    parser.add_argument('--wipe', action='store_true', help='Securely wipe original after encryption')
    
    args = parser.parse_args()
    
    try:
        encryptor = SecureEncryptor()
        
        if not args.password:
            import getpass
            args.password = getpass.getpass('Enter password: ')
        
        if os.path.isfile(args.path):
            if args.action == 'encrypt':
                output_path = encryptor.encrypt_file(args.path, args.password)
                print(f"File encrypted: {output_path}")
                if args.wipe:
                    encryptor.secure_wipe(args.path)
            else:
                output_path = encryptor.decrypt_file(args.path, args.password)
                print(f"File decrypted: {output_path}")
                
        elif os.path.isdir(args.path):
            if args.action == 'encrypt':
                output_path = encryptor.encrypt_directory(args.path, args.password)
                print(f"Directory encrypted: {output_path}")
                if args.wipe:
                    encryptor.secure_wipe(args.path)
            else:
                output_path = encryptor.decrypt_directory(args.path, args.password)
                print(f"Directory decrypted: {output_path}")
        
        else:
            print("Error: Path not found")
            sys.exit(1)
            
    except Exception as e:
        print(f"Error: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()