secure files manager for gitlab

glab_securefiles_manager.py

#!/usr/bin/env python3

import argparse
import os
import sys
import requests

GITLAB_URL = os.getenv("GITLAB_URL")
GITLAB_PRIVATE_TOKEN = os.getenv("GITLAB_PRIVATE_TOKEN")
GITLAB_PROJECT_ID = os.getenv("GITLAB_PROJECT_ID")


def check_env_vars():
    if not all([GITLAB_URL, GITLAB_PRIVATE_TOKEN, GITLAB_PROJECT_ID]):
        print("Error: Please set the following environment variables:")
        print("  - GITLAB_URL: Your GitLab instance URL (e.g., https://gitlab.com)")
        print(
            "  - GITLAB_PRIVATE_TOKEN: Your GitLab Personal Access Token with 'api' scope."
        )
        print("  - GITLAB_PROJECT_ID: Your GitLab project ID.")
        sys.exit(1)


def get_api_url(endpoint):
    return f"{GITLAB_URL}/api/v4/projects/{GITLAB_PROJECT_ID}{endpoint}"


def handle_api_error(response):
    if not response.ok:
        print(f"Error: API request failed with status code {response.status_code}")
        try:
            print(f"Response: {response.json()}")
        except ValueError:
            print(f"Response: {response.text}")
        sys.exit(1)


def list_secure_files():
    url = get_api_url("/secure_files?per_page=2500")
    headers = {"PRIVATE-TOKEN": GITLAB_PRIVATE_TOKEN}
    response = requests.get(url, headers=headers)
    handle_api_error(response)
    files = response.json()
    if not files:
        print("No secure files found in this project.")
        return

    print("Secure Files:")
    for f in files:
        print(f"  - ID: {f['id']}, Name: {f['name']}, Created: {f['created_at']}")


def download_secure_files(file_ids):
    print(f"Downloading {len(file_ids)} secure file(s)...")
    for file_id in file_ids:
        url = get_api_url(f"/secure_files/{file_id}/download")
        headers = {"PRIVATE-TOKEN": GITLAB_PRIVATE_TOKEN}

        with requests.get(url, headers=headers, stream=True) as r:
            handle_api_error(r)
            content_disposition = r.headers.get("content-disposition")
            if content_disposition:
                filename = content_disposition.split("filename=")[-1].strip('"')
            else:
                filename = f"secure_file_{file_id}"

            with open(filename, "wb") as f:
                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)
            print(f"  - Downloaded '{filename}'")


def upload_secure_files(file_paths):
    print(f"Uploading {len(file_paths)} file(s)...")
    for file_path in file_paths:
        if not os.path.exists(file_path):
            print(f"Error: File not found at '{file_path}'. Skipping.")
            continue

        url = get_api_url("/secure_files")
        headers = {"PRIVATE-TOKEN": GITLAB_PRIVATE_TOKEN}
        file_name = os.path.basename(file_path)

        with open(file_path, "rb") as f:
            files = {"file": (file_name, f, "application/octet-stream")}
            data = {"name": file_name}
            response = requests.post(url, headers=headers, data=data, files=files)

        if response.status_code == 201:
            print(f"  - Successfully uploaded '{file_name}'")
        else:
            handle_api_error(response)


def delete_secure_files(file_ids):
    print(f"Deleting {len(file_ids)} secure file(s)...")
    for file_id in file_ids:
        url = get_api_url(f"/secure_files/{file_id}")
        headers = {"PRIVATE-TOKEN": GITLAB_PRIVATE_TOKEN}
        response = requests.delete(url, headers=headers)
        if response.status_code == 204:
            print(f"  - Successfully deleted secure file with ID: {file_id}")
        else:
            handle_api_error(response)


def main():
    parser = argparse.ArgumentParser(
        description="A script to manage GitLab Secure Files via the API.",
        epilog="Make sure to set GITLAB_URL, GITLAB_PRIVATE_TOKEN, and GITLAB_PROJECT_ID environment variables.",
    )
    subparsers = parser.add_subparsers(
        dest="command", required=True, help="Available commands"
    )

    subparsers.add_parser("list", help="List all secure files.")

    download_parser = subparsers.add_parser(
        "download", help="Download one or more secure files by ID."
    )
    download_parser.add_argument(
        "file_ids", nargs="+", help="One or more secure file IDs to download."
    )

    upload_parser = subparsers.add_parser("upload", help="Upload one or more files.")
    upload_parser.add_argument(
        "file_paths", nargs="+", help="Paths to one or more files to upload."
    )

    delete_parser = subparsers.add_parser(
        "delete", help="Delete one or more secure files by ID."
    )
    delete_parser.add_argument(
        "file_ids", nargs="+", help="One or more secure file IDs to delete."
    )

    args = parser.parse_args()

    check_env_vars()

    if args.command == "list":
        list_secure_files()
    elif args.command == "download":
        download_secure_files(args.file_ids)
    elif args.command == "upload":
        upload_secure_files(args.file_paths)
    elif args.command == "delete":
        delete_secure_files(args.file_ids)


if __name__ == "__main__":
    main()