Christian Krakau-Louis
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import logging
|
|
import paramiko
|
|
from pathlib import Path
|
|
from app.config import settings
|
|
from app.celery_app import celery
|
|
from app.tasks.retry_config import BaseTaskWithRetry
|
|
from app.utils.filename_utils import get_unique_filename, sanitize_filename, extract_remote_path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
@celery.task(base=BaseTaskWithRetry)
|
|
def upload_to_sftp(file_path: str):
|
|
"""
|
|
Upload a file to an SFTP server.
|
|
"""
|
|
if not os.path.exists(file_path):
|
|
error_msg = f"File not found: {file_path}"
|
|
logger.error(error_msg)
|
|
raise FileNotFoundError(error_msg)
|
|
|
|
if not (settings.sftp_host and settings.sftp_port and settings.sftp_username):
|
|
logger.info("SFTP upload skipped: Missing configuration")
|
|
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()
|
|
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
|
|
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"Uploading {filename} to SFTP at {remote_path}")
|
|
sftp.put(file_path, remote_path)
|
|
logger.info(f"Successfully uploaded {filename} to SFTP at {remote_path}")
|
|
|
|
# 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:
|
|
pass
|
|
|
|
error_msg = f"Failed to upload {filename} to SFTP server: {str(e)}"
|
|
logger.error(error_msg)
|
|
raise Exception(error_msg)
|