簡單的實現文件上傳下載功能

簡單的實現文件上傳下載功能,其中用到了socketserver,可以實現多用戶上傳和下載文件

服務端代碼:

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Author: Xiaobai Lei
import socket
import subprocess
import os
import json
import hashlib
import struct
import socketserver

class MyServer(socketserver.BaseRequestHandler):
    """文件上傳下載服務端"""
    address_family = socket.AF_INET
    socket_type = socket.SOCK_STREAM    # tcp流類型
    allow_reuse_address = False
    max_packet_size = 8192
    coding = 'utf-8'
    request_queue_size = 5  # 隊列最長數量
    server_upload_dir = 'file_upload'
    server_download_dir = 'file_download'

    def server_close(self):
        self.request.close()

    def close_request(self, request):
        request.close()

    def handle(self):
        print('from client ', self.client_address)
        while True:
            try:
                head_struct = self.request.recv(4)
                if not head_struct:break

                head_len = struct.unpack('i', head_struct)[0]
                head_json = self.request.recv(head_len).decode(self.coding)
                head_dic = json.loads(head_json)

                action = head_dic['action']
                if hasattr(self, action):
                    func = getattr(self, action)
                    func(head_dic)
            except Exception:
                break

    def put(self, args):
        """實現文件上傳"""
        file_path = os.path.normpath(os.path.join(self.server_upload_dir, args['filename']))
        filesize = args['filesize']
        recv_size = 0  # 記錄讀取文件的大小
        file_md5 = hashlib.md5()    # 計算文件的md5值,保證上傳的完整性
        with open(file_path, "wb") as f:
            while recv_size < filesize:
                recv_data = self.request.recv(self.max_packet_size)
                file_md5.update(recv_data)
                f.write(recv_data)
                recv_size += len(recv_data)
        file_md5_code = file_md5.hexdigest()  # 獲取文件md5的值,字符串形式
        self.request.send(b"ok")
        recv_md5_code = self.request.recv(self.max_packet_size).decode(self.coding)    # 接收客戶端發來的md5值
        if recv_md5_code == file_md5_code:
            print("%s文件上傳成功,大小為%s字節!"%(args['filename'], filesize))
            self.request.send(b"1000")
        else:
            self.request.send(b"1001")

    def download(self, args):
        """實現文件下載"""
        file_path = os.path.normpath(os.path.join(self.server_download_dir, args['filename']))
        if not os.path.exists(file_path):
            print('file:%s is not exists' % file_path)
            return
        else:
            filesize = os.path.getsize(file_path)
        filesize_struct = struct.pack('i', filesize)
        # 發送服務端文件的大小
        self.request.send(filesize_struct)
        send_size = 0  # 記錄讀取文件的大小
        file_md5 = hashlib.md5()  # 計算文件的md5值,保證上傳的完整性
        # 接下來傳輸文件
        with open(file_path, "rb") as f:
            for line in f:
                file_md5.update(line)
                self.request.send(line)
                send_size += len(line)
        # 發送文件的md5,判斷文件一致性
        file_md5_code = file_md5.hexdigest()  # 獲取文件md5的值,字符串形式
        file_md5_client = self.request.recv(self.max_packet_size).decode(self.coding)
        if file_md5_client == file_md5_code:
            print("%s文件下載成功,大小為%s字節!" % (args['filename'], filesize))
            self.request.send(b"1000")
        else:
            self.request.send(b"1001")



if __name__ == '__main__':
    ip_port = ('127.0.0.1', 9999)
    server = socketserver.ThreadingTCPServer(ip_port, MyServer)
    server.serve_forever()

客戶端代碼:

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Author: Xiaobai Lei
import socket
import subprocess
import os
import json
import hashlib
import struct


class MyTcpClient:
    address_family = socket.AF_INET
    socket_type = socket.SOCK_STREAM
    allow_reuse_address = False
    max_packet_size = 8192
    coding = 'utf-8'
    request_queue_size = 5
    # client_download_dir = 'file_download'

    def __init__(self, server_address, connect=True):
        self.server_address = server_address
        self.socket = socket.socket(self.address_family,
                                    self.socket_type)
        if connect:
            try:
                self.client_connect()
            except:
                self.client_close()
                raise

    def client_connect(self):
        self.socket.connect(self.server_address)

    def client_close(self):
        self.socket.close()

    def run(self):
        while True:
            cmd = input(">>: ").strip()
            if not cmd:continue
            action,filename = cmd.split()
            if hasattr(self, action):
                func = getattr(self, action)
                func((action,filename))

    def put(self, args):
        """實現文件上傳"""
        action, filename = args[0],args[1]
        if not os.path.isfile(filename):
            print('file:%s is not exists' % filename)
            return
        else:
            filesize = os.path.getsize(filename)
        head_dic = {'action': action, 'filename': os.path.basename(filename), 'filesize': filesize}
        head_json = json.dumps(head_dic)
        head_json_bytes = head_json.encode('utf8')
        head_struct = struct.pack("i", len(head_json_bytes))
        # 先發送字典數據的pack字節
        self.socket.send(head_struct)
        # 接着發送字典數據
        self.socket.send(head_json_bytes)
        send_size = 0
        file_md5 = hashlib.md5()
        # 最后發送需要傳輸的文件
        with open(filename, 'rb') as f:
            for line in f:
                file_md5.update(line)
                self.socket.send(line)
                send_size += len(line)
        # 發送文件的md5,判斷文件一致性
        file_md5_code = file_md5.hexdigest()    # 獲取文件md5的值,字符串形式
        self.socket.recv(self.max_packet_size)
        self.socket.send(file_md5_code.encode(self.coding)) # 發送文件的md5值給服務端
        code = self.socket.recv(self.max_packet_size).decode(self.coding)
        if code == "1000":
            print("%s字節的%s文件上傳成功!"%(send_size, filename))
        else:
            print("上傳失敗,請重新上傳!")

    def download(self, args):
        """實現文件下載"""
        print(args)
        action, filename = args[0], args[1]
        head_dic = {'action': action, 'filename': os.path.basename(filename)}
        head_json = json.dumps(head_dic)
        head_json_bytes = head_json.encode('utf8')
        head_struct = struct.pack("i", len(head_json_bytes))
        # 先發送字典數據的pack字節
        self.socket.send(head_struct)
        # 接着發送字典數據
        self.socket.send(head_json_bytes)
        # 接收服務端要發送過來的文件大小數據
        filesize_struct = self.socket.recv(4)
        filesize = struct.unpack('i', filesize_struct)[0]
        recv_size = 0  # 記錄讀取文件的大小
        file_md5 = hashlib.md5()  # 計算文件的md5值,保證上傳的完整性
        with open(filename, "wb") as f:
            while recv_size < filesize:
                recv_data = self.socket.recv(self.max_packet_size)
                file_md5.update(recv_data)
                f.write(recv_data)
                recv_size += len(recv_data)
        file_md5_code = file_md5.hexdigest()  # 獲取文件md5的值,字符串形式
        self.socket.send(file_md5_code.encode(self.coding)) # 發送文件的md5值給服務端
        code = self.socket.recv(self.max_packet_size).decode(self.coding)
        if code == "1000":
            print("%s字節的%s文件下載成功!"%(recv_size, filename))
        else:
            print("下載失敗,請重新下載!")


if __name__ == '__main__':
    client = MyTcpClient(('127.0.0.1', 9999))
    client.run()