import docker
import argparse

# Initialize the Docker client
client = docker.from_env()

# Function to get the size of a Docker volume
def get_volume_size(container_id, volume_name):
    inspect_data = client.api.inspect_container(container_id)
    mounts = inspect_data['Mounts']
    for mount in mounts:
        if mount['Destination'] == volume_name:
            print(mount)
            return int(mount.get('Size',0))

# Function to shutdown a Docker container
def shutdown_container(container_id):
    client.api.stop(container_id)

# Define the threshold size in bytes
threshold_size = 1000000000  # 1GB

# Parse command-line arguments
parser = argparse.ArgumentParser(description='Check volume sizes of Docker containers and shutdown if size exceeds threshold.')
parser.add_argument('--dry-run', action='store_true', help='Perform a dry run without actually shutting down containers.')
args = parser.parse_args()

# Get a list of all running containers
containers = client.containers.list()

# Iterate through each container and check the size of its volumes
for container in containers:
    container_id = container.id
    print("Found container {}".format(container_id))
    inspect_data = client.api.inspect_container(container_id)
    mounts = inspect_data['Mounts']
    print("Found mounts {}".format(mounts))
    # Check the size of each volume and shutdown the container if any volume exceeds the threshold
    for mount in mounts:
        volume_name = mount['Destination']
        volume_size = get_volume_size(container_id, volume_name)
        if volume_size > threshold_size:
            if not args.dry_run:
                shutdown_container(container_id)
                print(f"Container {container_id} has been shutdown due to volume {volume_name} exceeding the threshold size.")
            else:
                print(f"Dry run: Container {container_id} would be shutdown due to volume {volume_name} exceeding the threshold size.")