cli.py

import argparse
import logging
import os
import sys
from backup import db_connect, backup, restore, compression, storage

def setup_logging(log_file):
    logging.basicConfig(
        filename=log_file,
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    logging.info(f"Logging initialized. Log file: {log_file}")


def load_env_file(env_file):
    if not env_file or not os.path.exists(env_file):
        return

    with open(env_file, 'r') as f:
        for raw_line in f:
            line = raw_line.strip()
            if not line or line.startswith('#') or '=' not in line:
                continue

            key, value = line.split('=', 1)
            key = key.strip()
            value = value.strip().strip('"').strip("'")

            if key and key not in os.environ:
                os.environ[key] = value


def build_config_from_env():
    db_port = os.getenv('DB_PORT')
    db_config = {
        'host': os.getenv('DB_HOST'),
        'port': int(db_port) if db_port and db_port.isdigit() else db_port,
        'user': os.getenv('DB_USER'),
        'password': os.getenv('DB_PASSWORD'),
        'database': os.getenv('DB_NAME'),
        'auth_database': os.getenv('DB_AUTH_DATABASE'),
    }

    aws_config = {
        'access_key': os.getenv('AWS_ACCESS_KEY'),
        'secret_key': os.getenv('AWS_SECRET_KEY'),
        'region': os.getenv('AWS_REGION'),
    }

    return {
        'database': db_config,
        'aws': aws_config,
    }


def validate_cloud_credentials(config, cloud_provider):
    if cloud_provider != 's3':
        raise ValueError(f"Unsupported cloud provider: {cloud_provider}. Only 's3' is supported.")

    aws_config = config.get('aws', {})
    if not aws_config.get('access_key') or not aws_config.get('secret_key') or not aws_config.get('region'):
        raise ValueError("AWS credentials are missing in environment variables.")


def validate_db_config(db_type, db_config):
    if db_type in ("mysql", "postgresql"):
        required_keys = ['host', 'port', 'user', 'password', 'database']
    elif db_type == "mongodb":
        required_keys = ['host', 'port', 'database']
    elif db_type == "sqlite":
        required_keys = ['database']
    else:
        raise ValueError(f"Unsupported database type: {db_type}")

    missing_keys = [key for key in required_keys if not db_config.get(key)]
    if missing_keys:
        missing = ", ".join(missing_keys)
        raise ValueError(f"Missing required database config for {db_type}: {missing}")

    if db_type == "mongodb":
        has_user = bool(db_config.get('user'))
        has_password = bool(db_config.get('password'))
        if has_user != has_password:
            raise ValueError("For mongodb, both 'user' and 'password' must be set together.")


def main():
    parser = argparse.ArgumentParser(description="Database Backup CLI Tool")
    
    # Main arguments
    parser.add_argument('operation', choices=['backup', 'restore'], help="Operation to perform: backup or restore")
    
    # Database connection args
    parser.add_argument('--db-type', required=True,choices=["mysql", "postgresql", "mongodb", "sqlite"], help="Database type (mysql, postgresql, mongodb, sqlite)")
    parser.add_argument('--env-file', default='.env', help="Optional path to an env file with secrets (default: .env)")
    
    # Backup args
    parser.add_argument('--output', help="Output backup file path")
    parser.add_argument('--compress', action='store_true', help="Compress the backup")
    
    # Restore args
    parser.add_argument('--backup-file', help="Backup file to restore")
    parser.add_argument('--s3-key', help="S3 object key to restore from. If omitted, latest object in bucket is used.")
    parser.add_argument('--s3-prefix', help="Optional S3 key prefix used when selecting latest backup.")
    parser.add_argument('--download-dir', default='.', help="Local directory to store backup downloaded from S3.")
    
    # Storage args
    parser.add_argument('--cloud', choices=['s3'], help="Cloud storage option")
    parser.add_argument('--bucket', help="Cloud storage bucket name")
    
    # Log file argument
    parser.add_argument('--log-file', help="Path to the log file", default='backup.log')
    
    # Parse the arguments
    args = parser.parse_args()

    setup_logging(args.log_file)

    logging.info(f"Starting {args.operation} operation.")

    load_env_file(args.env_file)
    config = build_config_from_env()
    db_config = config['database']
    
    try:
        validate_db_config(args.db_type, db_config)
    except ValueError as e:
        logging.error(str(e))
        print(str(e))
        return 1

    if args.operation == 'backup' and not args.output:
        logging.error("Error: --output is required for backup operations.")
        print("Error: --output is required for backup operations.")
        return 1
    
    if args.cloud:
        if not args.bucket:
            logging.error("Error: --bucket is required when --cloud is specified.")
            print("Error: --bucket is required when --cloud is specified.")
            return 1
        try:
            validate_cloud_credentials(config, args.cloud)
        except ValueError as e:
            logging.error(str(e))
            print(str(e))
            return 1
    try:
        if args.operation == 'backup':
            conn = db_connect.connect_to_db(args.db_type, db_config)
            if conn:
                try:
                    logging.info(f"Successfully connected to {args.db_type} database.")
                    backup_file = backup.full_backup(args.db_type, db_config['database'], db_config, args.output, logging)
                    logging.info(f"Backup created at {args.output}.")
                    
                    final_backup_path = backup_file
                    if args.compress:
                        final_backup_path = args.output + ".tar.gz"
                        compression.compress_backup(backup_file, final_backup_path, logger=logging)
                        logging.info(f"Backup compressed to {final_backup_path}")

                    if args.cloud:
                        storage.upload_to_cloud(args.cloud, final_backup_path, args.bucket, config, logger=logging)
                finally:
                    if hasattr(conn, "close"):
                        conn.close()
            else:
                logging.error(f"Failed to connect to {args.db_type} database.")
                return 1
        
        elif args.operation == 'restore':
            backup_path = args.backup_file

            if args.cloud:
                s3_key = args.s3_key
                if not s3_key:
                    s3_key = storage.get_latest_cloud_backup_key(
                        args.cloud,
                        args.bucket,
                        config,
                        logger=logging,
                        prefix=args.s3_prefix
                    )
                backup_path = storage.download_from_cloud(
                    args.cloud,
                    args.bucket,
                    s3_key,
                    args.download_dir,
                    config,
                    logger=logging
                )

            if not backup_path:
                logging.error("Error: provide --backup-file for local restore or --cloud s3 with --bucket for S3 restore.")
                print("Error: provide --backup-file for local restore or --cloud s3 with --bucket for S3 restore.")
                return 1

            restore.restore_backup(args.db_type, backup_path, db_config, logger=logging)
            logging.info(f"Database restored from {backup_path}.")
    
    except Exception as e:
        logging.error(f"An error occurred during the {args.operation} operation: {e}")
        print(f"An error occurred: {e}")
        return 1
    
    logging.info(f"{args.operation.capitalize()} operation completed.")
    return 0

if __name__ == '__main__':
    sys.exit(main())