jumpserver/apps/common/storage/jms_storage/sftp.py

110 lines
3.5 KiB
Python

# -*- coding: utf-8 -*-
import io
import os
import paramiko
from .base import ObjectStorage
class SFTPStorage(ObjectStorage):
def __init__(self, config):
self.sftp = None
self.sftp_host = config.get('SFTP_HOST', None)
self.sftp_port = int(config.get('SFTP_PORT', 22))
self.sftp_username = config.get('SFTP_USERNAME', '')
self.sftp_secret_type = config.get('STP_SECRET_TYPE', 'password')
self.sftp_password = config.get('SFTP_PASSWORD', '')
self.sftp_private_key = config.get('STP_PRIVATE_KEY', '')
self.sftp_passphrase = config.get('STP_PASSPHRASE', '')
self.sftp_root_path = config.get('SFTP_ROOT_PATH', '/tmp')
self.ssh = paramiko.SSHClient()
self.connect()
def connect(self):
self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
if self.sftp_secret_type == 'password':
self.ssh.connect(self.sftp_host, self.sftp_port, self.sftp_username, self.sftp_password)
elif self.sftp_secret_type == 'ssh_key':
pkey = paramiko.RSAKey.from_private_key(io.StringIO(self.sftp_private_key))
self.ssh.connect(self.sftp_host, self.sftp_port, self.sftp_username, pkey=pkey,
passphrase=self.sftp_passphrase)
self.sftp = self.ssh.open_sftp()
def confirm_connected(self):
try:
self.sftp.getcwd()
except Exception as e:
self.connect()
def upload(self, src, target):
local_file = src
remote_file = os.path.join(self.sftp_root_path, target)
try:
self.confirm_connected()
mode = os.stat(local_file).st_mode
remote_dir = os.path.dirname(remote_file)
if not self.exists(remote_dir):
self.sftp.mkdir(remote_dir)
self.sftp.put(local_file, remote_file)
self.sftp.chmod(remote_file, mode)
return True, None
except Exception as e:
return False, e
def download(self, src, target):
remote_file = src
local_file = target
self.confirm_connected()
try:
local_dir = os.path.dirname(local_file)
if not os.path.exists(local_dir):
os.makedirs(local_dir)
mode = self.sftp.stat(remote_file).st_mode
self.sftp.get(remote_file, local_file)
os.chmod(local_file, mode)
return True, None
except Exception as e:
return False, e
def delete(self, path):
path = os.path.join(self.sftp_root_path, path)
self.confirm_connected()
if not self.exists(path):
raise FileNotFoundError('File not exist error(%s)' % path)
try:
self.sftp.remove(path)
return True, None
except Exception as e:
return False, e
def check_dir_exist(self, d):
self.confirm_connected()
try:
self.sftp.stat(d)
return True
except Exception:
return False
def mkdir(self, dirs):
self.confirm_connected()
try:
if not self.exists(dirs):
self.sftp.mkdir(dirs)
return True, None
except Exception as e:
return False, e
def exists(self, target):
self.confirm_connected()
try:
self.sftp.stat(target)
return True
except:
return False
def close(self):
self.sftp.close()
self.ssh.close()