"""Import restriction wrapper"""

from typing import List
from diona.project.sandbox.wrappers.base import BuiltinWrapper


class ImportWrapper(BuiltinWrapper):
    """Restricts which modules can be imported"""
    
    def __init__(self, blocked_modules: List[str] = None):
        self._blocked_modules = blocked_modules or [
            'subprocess', 'socket', 'urllib', 'requests', 'http',
            'ftplib', 'smtplib', 'telnetlib', 'pty', 'tty', 'termios',
            'resource', 'signal', 'posix', 'nt', 'msvcrt', 'mmap', 'fcntl',
        ]
    
    def get_restricted_builtins(self) -> List[str]:
        return ['__import__']
    
    def get_wrapper_code(self) -> str:
        blocked = str(self._blocked_modules)
        return f'''_original_import = builtins.__import__
def _restricted_import(name, *args, **kwargs):
    if name in {blocked}:
        raise ImportError(f"Module '{{name}}' is not allowed in sandbox")
    return _original_import(name, *args, **kwargs)

builtins.__import__ = _restricted_import
'''
