gh-christianlouis-docuelevate/app/tasks/upload_to_sftp.py

#!/usr/bin/env python3

import logging
import os

import paramiko

from app.celery_app import celery
from app.config import settings
from app.tasks.retry_config import BaseTaskWithRetry
from app.utils import log_task_progress
from app.utils.filename_utils import extract_remote_path, get_unique_filename, sanitize_filename

logger = logging.getLogger(__name__)


@celery.task(base=BaseTaskWithRetry, bind=True)
def upload_to_sftp(self, file_path: str, file_id: int = None):
    """
    Upload a file to an SFTP server.

    Args:
        file_path: Path to the file to upload
        file_id: Optional file ID to associate with logs
    """
    task_id = self.request.id
    logger.info(f"[{task_id}] Starting SFTP upload: {file_path}")
    log_task_progress(
        task_id, "upload_to_sftp", "in_progress", f"Uploading to SFTP: {os.path.basename(file_path)}", file_id=file_id
    )

    if not os.path.exists(file_path):
        error_msg = f"File not found: {file_path}"
        logger.error(f"[{task_id}] {error_msg}")
        log_task_progress(task_id, "upload_to_sftp", "failure", error_msg, file_id=file_id)
        raise FileNotFoundError(error_msg)

    if not (settings.sftp_host and settings.sftp_port and settings.sftp_username):
        error_msg = "SFTP upload skipped: Missing configuration"
        logger.info(f"[{task_id}] {error_msg}")
        log_task_progress(task_id, "upload_to_sftp", "skipped", error_msg, file_id=file_id)
        return {"status": "Skipped", "reason": "SFTP settings not configured"}

    filename = os.path.basename(file_path)
    sanitized_filename = sanitize_filename(filename)

    # SSH client for SFTP connection
    ssh = paramiko.SSHClient()

    # Security: Host key verification
    # By default (sftp_disable_host_key_verification=False), we use RejectPolicy for security
    # Setting sftp_disable_host_key_verification=True disables verification (for testing only)
    if getattr(settings, "sftp_disable_host_key_verification", False):
        logger.warning(
            "SFTP host key verification is DISABLED - connections are vulnerable to MITM attacks. "
            "This should only be used in development/testing. For production, remove "
            "SFTP_DISABLE_HOST_KEY_VERIFICATION or set it to False and configure known_hosts."
        )
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # nosec B507 - Configurable, warns user  # noqa: S507
    else:
        # Use system known_hosts for host key verification (more secure)
        ssh.load_system_host_keys()
        ssh.set_missing_host_key_policy(paramiko.RejectPolicy())

    try:
        # Setup connection parameters
        connect_kwargs = {
            "hostname": settings.sftp_host,
            "port": settings.sftp_port,
            "username": settings.sftp_username,
        }

        # Check for authentication methods - use key if available, otherwise try password
        sftp_key_path = getattr(settings, "sftp_private_key", None)
        sftp_key_passphrase = getattr(settings, "sftp_private_key_passphrase", None)

        if sftp_key_path and os.path.exists(sftp_key_path):
            logger.info(f"Using SSH key authentication with key: {sftp_key_path}")
            connect_kwargs["key_filename"] = sftp_key_path
            if sftp_key_passphrase:
                connect_kwargs["passphrase"] = sftp_key_passphrase
        elif settings.sftp_password:
            logger.info("Using password authentication for SFTP")
            connect_kwargs["password"] = settings.sftp_password
        else:
            error_msg = "No authentication method available for SFTP (no key or password)"
            logger.error(error_msg)
            raise Exception(error_msg)

        # Connect to the server
        logger.info(f"Connecting to SFTP server at {settings.sftp_host}:{settings.sftp_port}")
        ssh.connect(**connect_kwargs)

        # Open SFTP session
        sftp = ssh.open_sftp()

        # Calculate remote path based on local file structure
        remote_base = settings.sftp_folder or ""
        remote_path = extract_remote_path(file_path, settings.workdir, remote_base)

        # Ensure the remote path starts with a slash if the base folder does
        if remote_base.startswith("/") and not remote_path.startswith("/"):
            remote_path = "/" + remote_path

        # Function to check if file exists in SFTP server
        def check_exists_in_sftp(path):
            try:
                sftp.stat(path)
                return True
            except FileNotFoundError:
                return False

        # Check for potential file collision and get a unique name if needed
        remote_path = get_unique_filename(remote_path, check_exists_in_sftp)

        # Create parent directories if needed
        remote_dir = os.path.dirname(remote_path)
        if remote_dir:
            try:
                # Try to create the full directory path
                current_dir = ""
                for dir_part in remote_dir.split("/"):
                    if not dir_part:
                        continue
                    current_dir += f"/{dir_part}"
                    try:
                        sftp.stat(current_dir)
                    except FileNotFoundError:
                        logger.info(f"Creating directory on SFTP server: {current_dir}")
                        sftp.mkdir(current_dir)
            except Exception as e:
                logger.warning(f"Failed to create directory structure {remote_dir}: {str(e)}")

        # Upload the file
        logger.info(f"[{task_id}] Uploading {filename} to SFTP at {remote_path}")
        sftp.put(file_path, remote_path)
        logger.info(f"[{task_id}] Successfully uploaded {filename} to SFTP at {remote_path}")
        log_task_progress(task_id, "upload_to_sftp", "success", f"Uploaded to SFTP: {filename}", file_id=file_id)

        # Close connections
        sftp.close()
        ssh.close()

        return {"status": "Completed", "file_path": file_path, "sftp_path": remote_path}

    except Exception as e:
        # Make sure connections are closed
        try:
            if "sftp" in locals():
                sftp.close()
            ssh.close()
        except Exception:  # noqa: S110
            pass  # Ignore errors during cleanup

        error_msg = f"Failed to upload {filename} to SFTP server: {str(e)}"
        logger.error(f"[{task_id}] {error_msg}")
        log_task_progress(task_id, "upload_to_sftp", "failure", error_msg, file_id=file_id)
        raise Exception(error_msg)