import resource
import os
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def set_resource_limits(cpu_time=30, memory_limit=1024 * 1024 * 512, open_files=1024, max_processes=512):
    """
    Sets resource limits to prevent resource exhaustion attacks.

    Args:
        cpu_time (int): Maximum CPU time in seconds.
        memory_limit (int): Maximum memory usage in bytes.
        open_files (int): Maximum number of open files.
        max_processes (int): Maximum number of processes a user can create.

    Raises:
        OSError: If setting a resource limit fails.  Logs the error and continues.
    """

    try:
        # Set CPU time limit
        resource.setrlimit(resource.RLIMIT_CPU, (cpu_time, cpu_time))
        logging.info(f"CPU time limit set to {cpu_time} seconds.")
    except OSError as e:
        logging.error(f"Failed to set CPU time limit: {e}")

    try:
        # Set memory limit
        resource.setrlimit(resource.RLIMIT_AS, (memory_limit, memory_limit)) # Address space limit
        logging.info(f"Memory limit set to {memory_limit / (1024 * 1024)} MB.")
    except OSError as e:
        logging.error(f"Failed to set memory limit: {e}")

    try:
        # Set open files limit
        resource.setrlimit(resource.RLIMIT_NOFILE, (open_files, open_files))
        logging.info(f"Open files limit set to {open_files}.")
    except OSError as e:
        logging.error(f"Failed to set open files limit: {e}")

    try:
        # Set max processes limit
        resource.setrlimit(resource.RLIMIT_NPROC, (max_processes, max_processes))
        logging.info(f"Max processes limit set to {max_processes}.")
    except OSError as e:
        logging.error(f"Failed to set max processes limit: {e}")

    try:
        # Limit core dump size to 0 to prevent large file creation
        resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
        logging.info("Core dump size limit set to 0.")
    except OSError as e:
        logging.error(f"Failed to set core dump size limit: {e}")

def apply_limits_to_process():
    """
    Applies resource limits to the current process.  Can be called at process start.
    """
    set_resource_limits()  # Use default limits or configure as needed.

if __name__ == '__main__':
    # Example usage (Demonstrates and Tests)
    apply_limits_to_process()

    # Verify limits (Prints the current limits)
    cpu_limit = resource.getrlimit(resource.RLIMIT_CPU)
    print(f"Current CPU Limit: {cpu_limit}")

    memory_limit = resource.getrlimit(resource.RLIMIT_AS)
    print(f"Current Memory Limit: {memory_limit}")

    open_files_limit = resource.getrlimit(resource.RLIMIT_NOFILE)
    print(f"Current Open Files Limit: {open_files_limit}")

    max_processes_limit = resource.getrlimit(resource.RLIMIT_NPROC)
    print(f"Current Max Processes Limit: {max_processes_limit}")

    core_dump_limit = resource.getrlimit(resource.RLIMIT_CORE)
    print(f"Current Core Dump Limit: {core_dump_limit}")

    # Attempt to exceed limits (Demonstration purposes - may crash the process)
    # NOTE:  These examples are for demonstration only.  Uncommenting them in a
    # production environment without proper handling can cause crashes.

    # CPU Limit Example: (Will raise TimeoutError after cpu_time limit)
    # import time
    # start_time = time.time()
    # while True:
    #     if time.time() - start_time > 60:  # Run longer than set cpu time.
    #         break
    #     pass
    # print("CPU intensive task completed (or timed out).")

    # Memory Limit Example: (Will raise MemoryError if memory_limit is reached)
    # try:
    #     large_list = bytearray(2 * 1024 * 1024 * 1024)  # Attempt to allocate 2GB
    #     print("Large list allocated successfully.")
    # except MemoryError as e:
    #     print(f"Memory allocation failed: {e}")