import threading
import time
import logging
import os

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

class ResourceExhaustionAttack:
    def __init__(self, attack_type, duration=10, cpu_load=80, memory_load=80, disk_load=80):
        self.attack_type = attack_type
        self.duration = duration
        self.cpu_load = cpu_load
        self.memory_load = memory_load
        self.disk_load = disk_load
        self.stop_event = threading.Event()

    def cpu_exhaustion(self):
        logging.info("Starting CPU exhaustion...")
        start_time = time.time()
        while not self.stop_event.is_set() and (time.time() - start_time) < self.duration:
            # Simulate CPU intensive task
            _ = [i**i for i in range(int(self.cpu_load * 1000))]  # Create CPU load based on percentage
            time.sleep(0.01)
        logging.info("CPU exhaustion stopped.")

    def memory_exhaustion(self):
        logging.info("Starting Memory exhaustion...")
        start_time = time.time()
        data = []
        while not self.stop_event.is_set() and (time.time() - start_time) < self.duration:
            try:
                # Allocate memory dynamically
                block_size = int(self.memory_load * 1024 * 1024 / 100)  # Allocate in MB based on percentage
                data.append('A' * block_size)
                time.sleep(0.01)
            except MemoryError:
                logging.error("MemoryError encountered during memory exhaustion.")
                break # Stop if we run out of memory

        # Clear allocated memory (important for graceful handling)
        del data
        logging.info("Memory exhaustion stopped.")

    def disk_exhaustion(self):
        logging.info("Starting Disk exhaustion...")
        start_time = time.time()
        file_path = "/tmp/exhaustion_file.txt"
        try:
            with open(file_path, 'wb') as f:
                while not self.stop_event.is_set() and (time.time() - start_time) < self.duration:
                    # Write large data blocks to disk
                    block_size = int(self.disk_load * 1024 * 1024 / 100) # Allocate in MB based on percentage
                    data = os.urandom(block_size)  # Generate random data
                    f.write(data)
                    f.flush() # Ensure data is written immediately
                    time.sleep(0.01)
        except IOError as e:
            logging.error(f"IOError encountered during disk exhaustion: {e}")
        finally:
            # Clean up the file
            try:
                os.remove(file_path)
                logging.info(f"Successfully removed {file_path}")
            except FileNotFoundError:
                logging.warning(f"File {file_path} not found during cleanup.")
            except Exception as e:
                logging.error(f"Error during file cleanup: {e}")

        logging.info("Disk exhaustion stopped.")


    def run(self):
        threads = []
        if "cpu" in self.attack_type:
            cpu_thread = threading.Thread(target=self.cpu_exhaustion)
            threads.append(cpu_thread)
            cpu_thread.start()
        if "memory" in self.attack_type:
            memory_thread = threading.Thread(target=self.memory_exhaustion)
            threads.append(memory_thread)
            memory_thread.start()
        if "disk" in self.attack_type:
            disk_thread = threading.Thread(target=self.disk_exhaustion)
            threads.append(disk_thread)
            disk_thread.start()

        # Wait for the specified duration
        time.sleep(self.duration)
        self.stop()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

    def stop(self):
        self.stop_event.set()
        logging.info("Stopping resource exhaustion attack...")


if __name__ == '__main__':
    # Example usage
    attack = ResourceExhaustionAttack(attack_type="cpu,memory,disk", duration=5, cpu_load=70, memory_load=60, disk_load=50)
    attack.run()