Python ftplib模块解决socket交互阻塞问题

使用ftplib模块操作FTP服务器时很容易阻塞，原因各种各样，有时还与ftp服务器有关。

下面的demo重载了ftplib模块，主要是使用了非阻塞的socket（self.sock.setblocking(0)），以及使用conn.recv接收数据的方式代替conn.makefile生成文件句柄接收数据，解决了阻塞问题，同时增加了一些功能。

#!usr/bin/env python
import os
import re
import time
import socket
import ftplib
import logging
import traceback
from ftplib import FTP, CRLF, Error, _SSLSocket


class FTP(FTP):
    blocksize = 8192

    def connect(self, host='', port=0, timeout=-999, source_address=None):
        if host != '':
            self.host = host
        if port > 0:
            self.port = port
        if timeout != -999:
            self.timeout = timeout
        if source_address is not None:
            self.source_address = source_address
        self.sock = socket.create_connection((self.host, self.port), self.timeout,
                                             source_address=self.source_address)
        self.sock.setblocking(0)  # 设置为非阻塞式连接
        self.af = self.sock.family
        self.file = self.sock.makefile('r', encoding=self.encoding)
        self.welcome = self.getresp()
        return self.welcome

    def retrbinary(self, cmd, callback, blocksize=8192, rest=None):
        self.voidcmd('TYPE I')
        with self.transfercmd(cmd, rest) as conn:
            while 1:
                try:
                    data = conn.recv(blocksize)
                except BlockingIOError:
                    time.sleep(0.05)
                    continue
                if not data:
                    break
                callback(data)
            if _SSLSocket is not None and isinstance(conn, _SSLSocket):
                conn.unwrap()
        return self.voidresp()

    def retrlines(self, cmd, callback=None):
        self.sendcmd('TYPE A')
        with self.transfercmd(cmd) as conn:
            result = b''
            while 1:
                try:
                    received = conn.recv(self.blocksize, socket.MSG_DONTWAIT)
                except BlockingIOError:
                    time.sleep(0.05)
                    continue
                if not received:
                    break
                result += received
            lines = result.decode(self.encoding).split('\n')
            for line in lines:
                if len(line) + 1 > self.maxline:
                    raise Error("got more than %d bytes" % self.maxline)
                if self.debugging > 2:
                    print('*retr*', repr(line + '\n'))
                if not line.strip():
                    break
                elif line[-1:] == '\r':
                    line = line[:-1]
                callback(line)
            if _SSLSocket is not None and isinstance(conn, _SSLSocket):
                conn.unwrap()
        return self.voidresp()

    def getline(self):
        start_time = time.time()
        while True:
            line = self.file.readline(self.maxline + 1)
            if line:
                break
            if time.time() - start_time > 3:
                raise EOFError
            time.sleep(0.05)
        if len(line) > self.maxline:
            raise Error("got more than %d bytes" % self.maxline)
        if self.debugging > 1:
            print('*get*', self.sanitize(line))
        if not line:
            raise EOFError
        if line[-2:] == CRLF:
            line = line[:-2]
        elif line[-1:] in CRLF:
            line = line[:-1]
        return line

    def ftp_login(self, host, port, username, password):  # 封装登录功能
        while 1:
            try:
                self.connect(host=host, port=int(port))
                self.login(user=username, passwd=password)
                return True  # 登录成功直接退出
            except (socket.error, socket.gaierror):
                logging.error('无法连接FTP主机%s\n%s' % (username, traceback.format_exc()))
            except ftplib.error_perm:
                logging.error('登录失败，请检查用户名和密码\n%s' % traceback.format_exc())
            except EOFError:
                self.exit()  # 登录失败断开连接
                logging.info('空网络IO导致登录失败，正在重试...')
                time.sleep(0.05)
                continue
            except:
                logging.error('登录失败，未知异常\n%s' % traceback.format_exc())
            self.exit()  # 登录失败断开连接
            return False

    def get_files(self, dirname, regex=''):  # 获取指定目录下文件，支持正则匹配
        files = []
        for file_path in self.nlst(dirname):
            try:
                self.cwd(file_path)  # 能进入的是目录
            except ftplib.error_perm:
                if not regex or re.match(r'%s' % regex, os.path.basename(file_path), re.S):  # 根据规则匹配文件
                    files.append(file_path)
        return files

    def isfile(self, file_path):  # 判断是不是文件
        try:
            self.cwd(file_path)
            return False  # 能进入的是目录
        except ftplib.error_perm as e:
            if 'Not a directory' in str(e):  # 提示不是一个目录
                return True
            else:
                return False  # 提示没有这个文件或目录

    def upload(self, upload_from, file_path):  # 上传文件
        try:
            with open(upload_from, 'rb') as file:
                self.storbinary('STOR %s' % file_path, file, self.blocksize)
        except:
            logging.error('文件[%s]上传失败\n%s' % (file_path, traceback.format_exc()))

    # 下载到mqt服务器
    def download(self, file_path, save_to):  # 下载文件
        try:
            with open(save_to, 'wb') as file:
                self.retrbinary('RETR %s' % file_path, file.write, self.blocksize)
            return True
        except:
            logging.error('文件[%s]下载失败\n%s' % (file_path, traceback.format_exc()))
            return False

    def exit(self):
        try:
            self.quit()
        except:
            pass


def upload_file(host, port, username, password, local_file_path, upload_to):
    try:
        ftp = FTP()
        ftp.set_debuglevel(0)  # 设置ftp日志级别
        if not ftp.ftp_login(host, port, username, password):
            return False
        ftp.upload(local_file_path, upload_to)
        ftp.exit()
        logging.info('文件[%s]上传成功' % local_file_path)
        return True
    except:
        ftp.exit()
        logging.error('文件[%s]上传失败\n%s' % (local_file_path, traceback.format_exc()))
        return False


def download_file(host, port, username, password, remote_file_path, save_to):
    try:
        ftp = FTP()
        ftp.set_debuglevel(0)  # 设置ftp日志级别
        if not ftp.ftp_login(host, port, username, password):
            return False
        if not ftp.isfile(remote_file_path):
            ftp.exit()
            logging.info('文件[%s]不存在' % remote_file_path)
            return False
        if not ftp.download(remote_file_path, save_to):
            ftp.exit()
            return False
        logging.info('文件[%s]下载成功' % remote_file_path)
        return True
    except:
        ftp.exit()
        logging.error('文件[%s]下载失败\n%s' % (remote_file_path, traceback.format_exc()))
        return False


def get_files(host, port, username, password, remote_path):
    try:
        ftp = FTP()
        ftp.set_debuglevel(0)  # 设置ftp日志级别
        if not ftp.ftp_login(host, port, username, password):
            return False
        # ftp.set_pasv(False)  # 切换主动和被动连接方式
        # 提取匹配的文件信息
        files = ftp.get_files(remote_path, r'^.*\.log$')  # 提取匹配的文件路径
        ftp.exit()
        return files
    except:
        ftp.exit()
        return False