socket_file.py

# -*- coding:utf-8 -*-
import socket
import struct
import os
import time
import hashlib
import threading


BUFFER_SIZE = 1024
HEAD_STRUCT = '128sIq32s'  # Structure of file head

def CalcMD5(filepath,block_size=2**20):
    with open(filepath,'rb') as f:
        md5obj = hashlib.md5()
        while 1:
            data = f.read(block_size)
            if not data:
                break
            md5obj.update(data)
        hash = md5obj.hexdigest()
        return hash

def connect_server(server_address):
    # Create a TCP/IP socket
    sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    # Connect the socket to the server
    try:
        # Connect
        sock.connect(server_address)
        print "Connecting to %s port %s" % server_address
    except Exception, e:
        print e

    return sock
        

def start_server(server_address):
    # Create a TCP socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(server_address)
    print "Starting server on %s port %s" % server_address
    sock.listen(1)
    return sock
    
def deal_filename(file):
    #处理文件路径中斜杠替换为当前系统的路径斜杠
    file = file.replace("\\", os.path.sep);
    file = file.replace("/", os.path.sep);
    l,r = os.path.split(file)
    if not os.path.exists(l):
        os.makedirs(l)
    return file

def send_file(sock,FILE_NAME,md5="n"):
    #发送文件函数
    FILE_SIZE = os.path.getsize(FILE_NAME)
    #计算MD5
    if not md5 == "n":
        print "Calculating MD5..."
        md5 = CalcMD5(FILE_NAME)
        print "Calculating success"

    # Need open again
    fr = open(FILE_NAME, 'rb')
    # Pack file info(file name and file size)
    file_head = struct.pack(HEAD_STRUCT, FILE_NAME, len(FILE_NAME), FILE_SIZE, md5)

    try:

        # Send file info\sock.send
        sock.send(file_head)
        print "Sending data..."
        time_start = time.time()
        while 1:
            file_data = fr.read(BUFFER_SIZE)
            if not file_data:
                break
            sock.send(file_data)

        time_end = time.time()
        fr.close()
        print "Send success!"
        print "MD5 : %s" % md5
        print "Cost %f seconds" % (time_end - time_start)
        print "Speed : %d MB/s"%(FILE_SIZE/1048576/(time_end - time_start))
        
        sock.close()
    except socket.errno, e:
        print "Socket error: %s" % str(e)
    except Exception, e:
        print "Other exception : %s" % str(e)
    finally:
        print "Closing connect"

def recv_file(client_socket,client_address):
    # Receive file info
    info_struct = struct.calcsize(HEAD_STRUCT)
    file_info = client_socket.recv(info_struct)

    if len(file_info) != 176:
        print file_info
        return 0

    file_name2, filename_size, file_size, md5_recv = struct.unpack(HEAD_STRUCT,file_info)
    file_name = file_name2[:filename_size].decode("utf-8")
    file_name = deal_filename(file_name)

    fw = open(file_name, 'wb')
    recv_size = 0
    print "Receiving data..."
    print "Data size:",file_size
    while 1:
        if file_size > BUFFER_SIZE:
            filedata = client_socket.recv(BUFFER_SIZE)
        else:
            filedata = client_socket.recv(file_size)
        if filedata:
            fw.write(filedata)
            file_size -=len(filedata)
        else:break
    fw.close()

    if str(md5_recv)[:1] == 'n':
        md5 = "No MD5 "

    else:
        print "Accept success!"
        print "Calculating MD5..."
        md5 = CalcMD5(file_name)
        print "  Recevie MD5 : %s" %md5_recv
        print "Calculate MD5 : %s" %md5

    return 1
def IterateFiles(directory):
    result = []
    if not os.path.isdir(directory):
        if directory == "":
            print "path can not end with \\!"
            print "Maybe path is not a directory or file"
            return 0
        result.append(directory)
        
    for root,dirs,files in os.walk(directory, topdown=True):
        for fl in files:
            file = fl.encode("utf-8")
            result.append(os.path.join(root,file))

    return result

def test_send():
    path = raw_input(">")
    l,r = os.path.split(path) 
    os.chdir(l)
    server_address = ("192.168.123.2", 998)
    files = IterateFiles(r)

    i = 1
    if not files:
        return 0
    for file in files :
        print "send %d  file"%i
        sock = connect_server(server_address)
        send_file(sock,file,"n")
        i +=1

    time.sleep(1)
    sock = connect_server(server_address)
    sock.send("break")
    sock.close()
    print "ok"
    

def test_recv():
    server_address = ("192.168.123.2", 998)
    sock = start_server(server_address)
    i = 1
    while i:
        client_socket,client_address = sock.accept()
        i = recv_file(client_socket,client_address)
        client_socket.close()
    sock.close()
    time.sleep(1)
    

if __name__ == '__main__':
    s = raw_input("do you want ser?")
    if s == "s":
        while 1:
            print "===================="
            test_recv()
            
    else:
        while 1:
            print "===================="
            test_send()