#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: JiaSongsong
# Date: 2014-12-21

import sys, os, socket, threading, logging, logging.handlers, time, stat, resource, md5
from optparse import OptionParser

FTP_LISTEN_IP       = '0.0.0.0'
FTP_LISTEN_PORT     = 21
MAX_CONNECTIONS     = 500
DEFAULT_HOME_DIR    = os.path.abspath('.')
ACCOUNT_INFO        = {
                        'ftp1'      :   'ftp1',
                        'anonymous' :   '',
                      }
IDLE_TIMEOUT        = 120
SOCKET_TIMEOUT      = 30
LOG_FILE_PATH       = 'ftp.py.log'
RW_BUFFER_SIZE      = 8192
ONLINE_COUNT        = 0

def add_online_count():
    global ONLINE_COUNT
    conn_lock = threading.Lock()
    conn_lock.acquire()
    ONLINE_COUNT += 1
    conn_lock.release()

def sub_online_count():
    global ONLINE_COUNT
    conn_lock = threading.Lock()
    conn_lock.acquire()
    ONLINE_COUNT -= 1
    conn_lock.release()

class FtpConnection(threading.Thread):
    def __init__(self, fd):
        threading.Thread.__init__(self)
        self.setDaemon(True)
        self.fd         = fd
        self.data_fd    = 0
        self.running    = True
        self.alive_time = time.time()
        self.options    = {'pasv': False, 'utf8': False}
        self.identified = False
        self.username   = ''
        self.home_dir   = DEFAULT_HOME_DIR
        self.curr_dir   = '/'
        self.file_pos   = 0
        self.handler    = dict()
        for method in dir(self):
            if method.startswith('handle_') and callable(getattr(self, method)):
                self.handler[method[7:]] = getattr(self, method)

    def run(self):
        add_online_count()
        try:
            self.say_welcome()
            while self.running:
                success, command, arg = self.recv()
                if self.idle_timeout():
                    logger.info('Ftp connection idle time too long, close it!')
                    break
                if not command or not success:
                    #self.send_msg(500, 'Failed')
                    continue
                command = command.upper()
                if not self.handler.has_key(command):
                    logger.error('Command Not Found: %s %s'%(command,arg))
                    self.send_msg(500, 'Command Not Found')
                    continue
                if self.options['utf8']:
                    arg = unicode(arg, 'utf8').encode(sys.getfilesystemencoding())
                logger.info('<<< %s %s'%(command,arg))
                self.keep_alive()
                if not self.identified and command not in ('USER','PASS'):
                    logger.error('Please login with USER and PASS.')
                    self.send_msg(530, 'Please login with USER and PASS.')
                try:
                    self.handler[command](arg)
                except OSError, e:
                    logger.error('Handle command error: %s'%e, exc_info=True)
                    self.send_msg(500, 'Permission denied')
            self.say_bye()
        except Exception, e:
            self.running = False
            logger.error('Error: %s'%e, exc_info=True)
        finally:
            self.fd.close()
        logger.info('FTP connnection <%r> done.'%self.fd)
        sub_online_count()

    def send_msg(self, code, msg, joiner_char=' '):
        if self.options['utf8']:
            msg = unicode(msg, sys.getfilesystemencoding()).encode('utf8')
        message = str(code) + joiner_char + msg + '\r\n'
        self.fd.send(message)
        logger.info('>>> %d%s%s'%(code,joiner_char,msg))

    def recv(self):
        '''returns 3 tuples, success, command, arg'''
        try:
            success, buf, command, arg = True, '', '', ''
            while True:
                data = self.fd.recv(RW_BUFFER_SIZE)
                if not data or data <= 0:
                    self.running = False
                    success = False
                    break
                buf += data
                if buf[-2:] == '\r\n': break
            split = buf.find(' ')
            command, arg = (buf[:split], buf[split + 1:].strip()) if split != -1 else (buf.strip(), '')
        except socket.error:
            errno, errstr = sys.exc_info()[:2]
            if errno != socket.timeout:
                logger.error('Receive data error: %d-%s'%(errno,errstr), exc_info=True)
                self.running = False
                success = False
        return success, command, arg

    def say_welcome(self):
        self.send_msg(220, 'Welcome to Python FtpServer!')

    def say_bye(self):
        self.handle_BYE('')

    def keep_alive(self):
        self.alive_time = time.time()

    def idle_timeout(self):
        return (time.time()- self.alive_time) >= IDLE_TIMEOUT

    def data_connect(self):
        '''Establish data connection'''
        if self.data_fd == 0:
            self.send_msg(500, 'no data connection')
            logger.info('No data connection')
            return False
        elif self.options['pasv']:
            fd, addr = self.data_fd.accept()
            self.data_fd.close()
            self.data_fd = fd
            logger.info('Data connection is established.')
        else:
            try:
                self.data_fd.connect((self.data_host, self.data_port))
                logger.info('Data connection is established.')
            except:
                self.send_msg(500, 'Failed to connect')
                logger.info('Failed to connect [%s:%d]'%(self.data_host, self.data_port))
                return False
        return True

    def close_data_fd(self):
        self.data_fd.close()
        self.data_fd = 0

    def parse_path(self, path):
        if path == '': path = '.'
        if path[0] != '/': path = self.curr_dir + os.sep + path
        logger.info('parse_path ' + path)
        split_path = os.path.normpath(path).replace('\\', '/').split('/')
        remote = ''
        local = self.home_dir # lock home directory
        for item in split_path:
            if item.startswith('..') or item == '': continue # ignore parent directory
            remote += '/' + item
            local  += os.sep + item
        if remote == '': remote = '/'
        #logger.info(split_path)
        logger.info('remote: %s, local: %s' % (remote, local))
        return remote, local

    def get_size(self, start_path):
        '''Get file or directory size'''
        total_size = 0
        if os.path.isfile(start_path):
            total_size = os.path.getsize(start_path)
        else:
            for dirpath, dirnames, filenames in os.walk(start_path):
                for f in filenames:
                    filepath = os.path.join(dirpath, f)
                    total_size += os.path.getsize(filepath)
        return total_size

    # Command Handlers
    def handle_USER(self, arg):
        if arg in ACCOUNT_INFO:
            self.username = arg
            if self.username == 'anonymous':
                self.send_msg(230, 'Login successful.')
                self.identified = True
                self.keep_alive()
            else:
                self.send_msg(331, 'User name okay, need password.')
        else:
            self.send_msg(500, 'Invalid User.')

    def handle_PASS(self, arg):
        if arg == ACCOUNT_INFO[self.username]:
            self.send_msg(230, 'Login successful.')
            self.identified = True
            self.keep_alive()
        else:
            self.send_msg(530, 'Password is not corrected')
            self.identified = False

    def handle_QUIT(self, arg):
        self.handle_BYE(arg)
    def handle_BYE(self, arg):
        self.running = False
        #self.send_msg(200, 'Bye')

    def handle_XPWD(self, arg):
        self.handle_PWD(arg)
    def handle_PWD(self, arg):
        remote, local = self.parse_path(self.curr_dir)
        self.send_msg(257, '"%s" is current directory'%remote)

    def handle_CDUP(self, arg):
        self.handle_CWD('..')
    def handle_CWD(self, arg):
        remote, local = self.parse_path(arg)
        try:
            os.listdir(local)
            self.curr_dir = remote
            self.send_msg(250, 'CWD command successful.')
        except Exception, e:
            logger.error('in cwd', exc_info=True)
            self.send_msg(500, 'Change directory failed!')

    def handle_SIZE(self, arg):
        remote, local = self.parse_path(arg)
        self.send_msg(213, str(self.get_size(local)))

    def handle_SYST(self, arg):
        self.send_msg(215, 'UNIX')

    def handle_APPE(self, arg):
        self.handle_STOR(arg, 'ab')
    def handle_STOR(self, arg, open_flag='wb'):
        remote, local = self.parse_path(arg)
        if not self.data_connect(): return
        self.send_msg(125, 'Data connection already open; transfer starting.')
        with open(local, open_flag) as f:
            f.seek(self.file_pos)
            while True:
                self.keep_alive()
                data = self.data_fd.recv(RW_BUFFER_SIZE)
                if len(data) == 0: break
                f.write(data)
        self.send_msg(226, 'Closing data connection.')
        self.close_data_fd()

    def handle_RETR(self, arg):
        remote, local = self.parse_path(arg)
        if not self.data_connect(): return
        self.send_msg(125, 'Data connection already open; transfer starting.')
        with open(local, 'rb') as f:
            f.seek(self.file_pos)
            while True:
                self.keep_alive()
                data = f.read(RW_BUFFER_SIZE)
                if not data: break
                self.data_fd.sendall(data)
        self.send_msg(226, 'Closing data connection.')
        self.close_data_fd()

    def handle_TYPE(self, arg):
        mode = 'Binary' if arg == 'I' else 'ASCII'
        self.send_msg(200, 'Switching to %s mode.'%mode)

    def handle_RNFR(self, arg):
        remote, local = self.parse_path(arg)
        self.rename_tmp_path = local
        self.send_msg(350, 'rename from ' + remote)

    def handle_RNTO(self, arg):
        remote, local = self.parse_path(arg)
        os.rename(self.rename_tmp_path, local)
        self.send_msg(250, 'rename to ' + remote)

    def handle_XMKD(self, arg):
        self.handle_MKD(arg)
    def handle_MKD(self, arg):
        remote, local = self.parse_path(arg)
        if os.path.exists(local):
            self.send_msg(500, 'Folder is already existed')
            return
        os.mkdir(local)
        self.send_msg(257, 'OK')

    def handle_XRMD(self, arg):
        self.handle_RMD(arg)
    def handle_RMD(self, arg):
        remote, local = self.parse_path(arg)
        if not os.path.exists(local):
            self.send_msg(500, 'Folder is not existed')
            return
        os.rmdir(local)
        self.send_msg(250, 'OK')

    def handle_NLST(self, arg):
        if not arg: arg = self.curr_dir
        if not self.data_connect(): return
        self.send_msg(150, 'File status okay; about to open data connection.')
        remote, local = self.parse_path(arg)
        filelist = ''
        for filename in os.listdir(local):
            filelist += arg + os.sep + filename + '\r\n'
        #logger.info('filelist:\n%s'%filelist)
        self.data_fd.sendall(filelist)
        self.send_msg(226, 'Closing data connection.')
        self.close_data_fd()

    def handle_LIST(self, arg):
        if not arg: arg = self.curr_dir
        if not self.data_connect(): return
        self.send_msg(150, 'File status okay; about to open data connection.')
        template = '%s%s%s------- %04u %8s %8s %8lu %s %s\r\n'
        remote, local = self.parse_path(arg)
        filelist = ''
        for filename in os.listdir(local):
            path = local + os.sep + filename
            if os.path.isfile(path) or os.path.isdir(path): # ignores link or block file
                status = os.stat(path)
                msg = template % (
                    'd' if os.path.isdir(path) else '-',
                    'r', 'w', 1, '0', '0',
                    status[stat.ST_SIZE],
                    time.strftime('%b %d  %Y', time.localtime(status[stat.ST_MTIME])),
                    filename)
                if self.options['utf8']: msg = unicode(msg, sys.getfilesystemencoding()).encode('utf8')
                filelist += msg
        self.data_fd.sendall(filelist)
        self.send_msg(226, 'Closing data connection.')
        self.close_data_fd()

    def handle_PASV(self, arg):
        self.options['pasv'] = True
        try:
            localhost = self.fd.getsockname()[0]
            self.data_fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.data_fd.bind((localhost, 0))
            self.data_fd.listen(1)
            ip, port = self.data_fd.getsockname()
            self.send_msg(227, 'Enter Passive Mode (%s,%u,%u).' %
                    (','.join(ip.split('.')), (port >> 8 & 0xff), (port & 0xff)))
        except Exception, e:
            logger.error('in pasv', exc_info=True)
            self.send_msg(500, 'Passive mode failed')

    def handle_PORT(self, arg):
        try:
            if self.data_fd:
                self.data_fd.close()
            t = arg.split(',')
            self.data_host = '.'.join(t[:4])
            self.data_port = int(t[4]) << 8 | int(t[5])
            self.data_fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        except:
            self.send_msg(500, 'PORT failed')
        self.send_msg(200, 'OK')

    def handle_DELE(self, arg):
        remote, local = self.parse_path(arg)
        if not os.path.exists(local):
            self.send_msg(450, 'File not exist')
            return
        os.remove(local)
        self.send_msg(250, 'File deleted')

    def handle_OPTS(self, arg):
        if arg.upper() == 'UTF8 ON':
            self.options['utf8'] = True
            self.send_msg(200, 'OK')
        elif arg.upper() == 'UTF8 OFF':
            self.options['utf8'] = False
            self.send_msg(200, 'OK')
        else:
            self.send_msg(500, 'Invalid argument')

    def handle_FEAT(self, arg):
        features = 'Features:\r\nMDTM\r\nPASV\r\nREST STREAM\r\nSIZE\r\nUTF8\r\nXMD5\r\n211 End'
        self.send_msg(211, features, '-')

    def handle_REST(self, arg):
        self.file_pos = int(arg)
        self.send_msg(350, 'Offset is set to %d.'%self.file_pos)

    def handle_MDTM(self, arg):
        remote, local = self.parse_path(arg)
        if os.path.exists(local):
            self.send_msg(213, time.strftime('%Y%m%d%I%M%S', time.localtime(os.path.getmtime(newpath))))
        else:
            self.send_msg(550, 'File is not exist')

    def _get_file_md5sum(self, path):
        hash_md5 = md5.md5()
        with open(path, "rb") as f:
            for chunk in iter(lambda: f.read(RW_BUFFER_SIZE), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()

    def handle_XMD5(self, arg):
        remote, local = self.parse_path(arg)
        if os.path.exists(local):
            self.send_msg(200, self._get_file_md5sum(local))
        else:
            self.send_msg(550, 'File is not exist')

def server_listen():
    listen_fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listen_fd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listen_fd.bind((FTP_LISTEN_IP, FTP_LISTEN_PORT))
    listen_fd.listen(1024)
    print 'FtpServer is listening on %s:%d'%(FTP_LISTEN_IP, FTP_LISTEN_PORT)

    # Set the default timeout in seconds (float) for new socket objects.
    socket.setdefaulttimeout(SOCKET_TIMEOUT)
    while True:
        if ONLINE_COUNT > MAX_CONNECTIONS:
            self.message(500, "Too many connections!")
            logger.error('Ftp connecions is greater than MAX_CONNECTIONS, reject new connection!')
        else:
            conn_fd, remote_addr = listen_fd.accept()
            logger.info('Connection from %r, online user count [%d]'%(remote_addr,ONLINE_COUNT))
            conn = FtpConnection(conn_fd)
            conn.start()

def get_logger(handler = logging.StreamHandler()):
    logger = logging.getLogger()
    formatter = logging.Formatter('[%(asctime)s][Thread:%(thread)d][%(levelname)s] %(message)s -- %(filename)s:%(funcName)s()@%(lineno)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.NOTSET)
    return logger

def daemonize(stdin='/dev/null', stdout='/dev/null', stderr='/dev/null'):
    '''becomes a daemon'''
    # Fork a child process so the parent can exit.
    try:
        pid = os.fork()
        if pid > 0: sys.exit(0)
    except OSError, e:
        sys.stderr.write("fork #1 failed\n")
        sys.exit(1)

    os.chdir('/')
    os.umask(0)
    os.setsid()
    # Fork a second child and exit immediately to prevent zombies.
    try:
        pid = os.fork()
        if pid > 0: sys.exit(0)
    except OSError, e:
        sys.stderr.write("fork #2 failed\n")
        sys.exit(1)

    for f in sys.stdout, sys.stderr: f.flush()
    si = file(stdin, 'r')
    so = file(stdout, 'a+')
    se = file(stderr, 'a+', 0)
    os.dup2(si.fileno(), sys.stdin.fileno())  # 0
    os.dup2(so.fileno(), sys.stdout.fileno()) # 1
    os.dup2(se.fileno(), sys.stderr.fileno()) # 2

def param_handler():
    global DEFAULT_HOME_DIR, FTP_LISTEN_IP, FTP_LISTEN_PORT, LOG_FILE_PATH, logger
    parser = OptionParser()
    parser.add_option('-d','--daemon',action="store_true",help='become a daemon')
    parser.add_option('-i','--ip',help='listen ip')
    parser.add_option('-p','--port',type=int,help='listen port')
    parser.add_option('-l','--logfile',help='log file path, default is "./ftp.py.log"')
    parser.add_option('-o','--stdout',action="store_true",help='output log to stdout, by default, it outputs to a log file')
    parser.add_option('-H','--home',help='ftp home directory, default is current directory')
    (options,args) = parser.parse_args()
    if options.daemon:  daemonize()
    if options.ip:      FTP_LISTEN_IP = options.ip
    if options.port:    FTP_LISTEN_PORT = options.port
    if options.home:    DEFAULT_HOME_DIR = os.path.abspath(options.home)
    if options.logfile:
        LOG_FILE_PATH = options.logfile
    logger = get_logger(logging.FileHandler(LOG_FILE_PATH))
    if options.stdout:  logger = get_logger()

def main():
    param_handler()

    # Set the maximum number of open file descriptors for the current process
    resource.setrlimit(resource.RLIMIT_NOFILE, (65535,65535))

    try:
        server_listen()
    except KeyboardInterrupt:
        print '^C received, Shutting Down Ftp Server'

if __name__ == '__main__':
    main()