cli.py

#!/usr/bin/env python3
"""
Artifact Repository CLI Management Tool
Command-line interface for managing the centralized artifact repository
"""

import click
import os
import sys
import json
import requests
from pathlib import Path
from typing import Optional
import yaml

class ArtifactRegistryClient:
    def __init__(self, base_url: str = "http://localhost:5000", token: Optional[str] = None):
        self.base_url = base_url.rstrip('/')
        self.token = token
        self.session = requests.Session()
        if token:
            self.session.headers.update({'Authorization': f'Bearer {token}'})
    
    def login(self, username: str, password: str) -> str:
        """Login and get access token."""
        response = self.session.post(f"{self.base_url}/api/auth/login", 
                                   json={'username': username, 'password': password})
        if response.status_code == 200:
            data = response.json()
            self.token = data['access_token']
            self.session.headers.update({'Authorization': f'Bearer {self.token}'})
            return self.token
        else:
            raise Exception(f"Login failed: {response.json().get('error', 'Unknown error')}")
    
    def upload_artifact(self, file_path: str, name: str, version: str, 
                       artifact_type: str, org: str, repo: str, 
                       description: str = "", tags: list = None) -> dict:
        """Upload an artifact to the repository."""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        with open(file_path, 'rb') as f:
            files = {'file': f}
            data = {
                'name': name,
                'version': version,
                'description': description,
                'tags': json.dumps(tags or [])
            }
            
            response = self.session.post(
                f"{self.base_url}/api/organizations/{org}/repositories/{repo}/artifacts",
                files=files,
                data=data
            )
        
        if response.status_code == 201:
            return response.json()
        else:
            raise Exception(f"Upload failed: {response.json().get('error', 'Unknown error')}")
    
    def create_organization(self, name: str, display_name: str = "", 
                          description: str = "", is_public: bool = False) -> dict:
        """Create a new organization."""
        data = {
            'name': name,
            'display_name': display_name or name,
            'description': description,
            'is_public': is_public
        }
        
        response = self.session.post(f"{self.base_url}/api/organizations", json=data)
        if response.status_code == 201:
            return response.json()
        else:
            raise Exception(f"Failed to create organization: {response.json().get('error', 'Unknown error')}")
    
    def create_repository(self, org: str, name: str, artifact_type: str,
                         display_name: str = "", description: str = "", 
                         is_public: bool = False) -> dict:
        """Create a new repository."""
        data = {
            'name': name,
            'display_name': display_name or name,
            'description': description,
            'artifact_type': artifact_type,
            'is_public': is_public
        }
        
        response = self.session.post(
            f"{self.base_url}/api/organizations/{org}/repositories", 
            json=data
        )
        if response.status_code == 201:
            return response.json()
        else:
            raise Exception(f"Failed to create repository: {response.json().get('error', 'Unknown error')}")
    
    def list_artifacts(self, org: str, repo: str, search: str = "", page: int = 1) -> dict:
        """List artifacts in a repository."""
        params = {'page': page}
        if search:
            params['search'] = search
            
        response = self.session.get(
            f"{self.base_url}/api/organizations/{org}/repositories/{repo}/artifacts",
            params=params
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Failed to list artifacts: {response.json().get('error', 'Unknown error')}")
    
    def get_stats(self) -> dict:
        """Get repository statistics."""
        response = self.session.get(f"{self.base_url}/api/stats")
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Failed to get stats: {response.json().get('error', 'Unknown error')}")

@click.group()
@click.option('--url', default='http://localhost:5000', help='Registry URL')
@click.option('--token', help='Access token')
@click.pass_context
def cli(ctx, url, token):
    """Centralized Artifact Repository Manager CLI"""
    ctx.ensure_object(dict)
    ctx.obj['client'] = ArtifactRegistryClient(url, token)

@cli.command()
@click.option('--username', prompt=True, help='Username')
@click.option('--password', prompt=True, hide_input=True, help='Password')
@click.pass_context
def login(ctx, username, password):
    """Login to the registry"""
    try:
        token = ctx.obj['client'].login(username, password)
        click.echo(f"Login successful! Token: {token}")
        
        # Save token to config file
        config_dir = Path.home() / '.artifact-registry'
        config_dir.mkdir(exist_ok=True)
        config_file = config_dir / 'config.json'
        
        config = {}
        if config_file.exists():
            with open(config_file) as f:
                config = json.load(f)
        
        config['token'] = token
        config['url'] = ctx.obj['client'].base_url
        
        with open(config_file, 'w') as f:
            json.dump(config, f, indent=2)
        
        click.echo(f"Token saved to {config_file}")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

@cli.command()
@click.argument('name')
@click.option('--display-name', help='Display name')
@click.option('--description', help='Description')
@click.option('--public', is_flag=True, help='Make organization public')
@click.pass_context
def create_org(ctx, name, display_name, description, public):
    """Create a new organization"""
    try:
        result = ctx.obj['client'].create_organization(
            name, display_name or '', description or '', public
        )
        click.echo(f"Organization '{result['name']}' created successfully!")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

@cli.command()
@click.argument('org')
@click.argument('name')
@click.argument('type', type=click.Choice([
    'docker', 'apt', 'npm', 'python', 'maven', 'nuget', 'helm',
    'generic', 'ai_model', 'dataset', 'firmware', 'documentation'
]))
@click.option('--display-name', help='Display name')
@click.option('--description', help='Description')
@click.option('--public', is_flag=True, help='Make repository public')
@click.pass_context
def create_repo(ctx, org, name, type, display_name, description, public):
    """Create a new repository"""
    try:
        result = ctx.obj['client'].create_repository(
            org, name, type, display_name or '', description or '', public
        )
        click.echo(f"Repository '{result['name']}' created successfully!")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

@cli.command()
@click.argument('file_path', type=click.Path(exists=True))
@click.argument('org')
@click.argument('repo')
@click.option('--name', help='Artifact name (defaults to filename)')
@click.option('--version', default='1.0.0', help='Version')
@click.option('--description', help='Description')
@click.option('--tags', help='Comma-separated tags')
@click.pass_context
def upload(ctx, file_path, org, repo, name, version, description, tags):
    """Upload an artifact"""
    try:
        if not name:
            name = Path(file_path).stem
        
        tag_list = []
        if tags:
            tag_list = [tag.strip() for tag in tags.split(',')]
        
        # Detect artifact type based on file extension
        file_ext = Path(file_path).suffix.lower()
        artifact_type = 'generic'  # default
        
        for atype, config in {
            'docker': ['.tar', '.tar.gz', '.tar.xz'],
            'apt': ['.deb'],
            'npm': ['.tgz'],
            'python': ['.whl', '.zip'],
            'maven': ['.jar', '.war', '.ear'],
            'ai_model': ['.pkl', '.h5', '.pb', '.onnx', '.pt', '.pth'],
            'dataset': ['.csv', '.json', '.parquet'],
            'firmware': ['.bin', '.hex', '.fw', '.rom'],
            'documentation': ['.pdf', '.md', '.html']
        }.items():
            if file_ext in config:
                artifact_type = atype
                break
        
        result = ctx.obj['client'].upload_artifact(
            file_path, name, version, artifact_type, org, repo, 
            description or '', tag_list
        )
        
        click.echo(f"Artifact '{result['name']}' v{result['version']} uploaded successfully!")
        click.echo(f"UUID: {result['uuid']}")
        click.echo(f"Size: {result['file_size']} bytes")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

@cli.command()
@click.argument('org')
@click.argument('repo')
@click.option('--search', help='Search term')
@click.option('--page', default=1, help='Page number')
@click.pass_context
def list_artifacts(ctx, org, repo, search, page):
    """List artifacts in a repository"""
    try:
        result = ctx.obj['client'].list_artifacts(org, repo, search or '', page)
        
        if not result['artifacts']:
            click.echo("No artifacts found.")
            return
        
        click.echo(f"\nArtifacts in {org}/{repo}:")
        click.echo("-" * 80)
        
        for artifact in result['artifacts']:
            click.echo(f"Name: {artifact['name']}")
            click.echo(f"Version: {artifact['version']}")
            click.echo(f"UUID: {artifact['uuid']}")
            click.echo(f"Size: {artifact['file_size']} bytes")
            click.echo(f"Downloads: {artifact['download_count']}")
            click.echo(f"Created: {artifact['created_at']}")
            click.echo(f"Owner: {artifact['owner']}")
            if artifact.get('tags'):
                click.echo(f"Tags: {', '.join(artifact['tags'])}")
            click.echo("-" * 80)
        
        pagination = result['pagination']
        click.echo(f"\nPage {pagination['page']} of {pagination['pages']} "
                  f"({pagination['total']} total artifacts)")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

@cli.command()
@click.pass_context
def stats(ctx):
    """Show repository statistics"""
    try:
        result = ctx.obj['client'].get_stats()
        
        click.echo("\n📊 Repository Statistics")
        click.echo("=" * 50)
        click.echo(f"Total Artifacts: {result['total_artifacts']}")
        click.echo(f"Total Organizations: {result['total_organizations']}")
        click.echo(f"Total Repositories: {result['total_repositories']}")
        click.echo(f"Total Downloads: {result['total_downloads']}")
        click.echo(f"Storage Usage: {result['storage_usage'] / (1024*1024*1024):.2f} GB")
        
        if result['artifact_types']:
            click.echo("\n📦 Artifact Types:")
            for atype, info in result['artifact_types'].items():
                click.echo(f"  {info['name']}: {info['count']}")
        
        if result['recent_uploads']:
            click.echo("\n🔄 Recent Uploads:")
            for artifact in result['recent_uploads'][:5]:
                click.echo(f"  {artifact['name']} v{artifact['version']} "
                          f"({artifact['type']}) by {artifact['owner']}")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

@cli.command()
@click.argument('config_file', type=click.Path(exists=True))
@click.pass_context
def bulk_upload(ctx, config_file):
    """Bulk upload artifacts from configuration file"""
    try:
        with open(config_file) as f:
            if config_file.endswith('.yaml') or config_file.endswith('.yml'):
                config = yaml.safe_load(f)
            else:
                config = json.load(f)
        
        for item in config.get('artifacts', []):
            click.echo(f"Uploading {item['file_path']}...")
            
            result = ctx.obj['client'].upload_artifact(
                item['file_path'],
                item['name'],
                item['version'],
                item['type'],
                item['organization'],
                item['repository'],
                item.get('description', ''),
                item.get('tags', [])
            )
            
            click.echo(f"✅ {result['name']} v{result['version']} uploaded")
        
        click.echo(f"\n🎉 Bulk upload completed!")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

if __name__ == '__main__':
    # Try to load saved config
    config_file = Path.home() / '.artifact-registry' / 'config.json'
    if config_file.exists():
        try:
            with open(config_file) as f:
                config = json.load(f)
            
            # Update CLI context with saved config
            import sys
            if '--url' not in sys.argv and config.get('url'):
                sys.argv.extend(['--url', config['url']])
            if '--token' not in sys.argv and config.get('token'):
                sys.argv.extend(['--token', config['token']])
        except:
            pass
    
    cli()