R1gelX`Blog

22 object(s)
 

初学Socket笔记

前段时间做题遇到了socket,于是现在写一下脚本来总结一下学到的一点基础的socket知识

基于socket写了一个客户端与服务端交互的脚本,可用于上传下载文件命令执行

客户端

# coding=utf-8
import os
from posixpath import split
import socket
import time
import hashlib
#client
host = '101.132.238.43'
port = 5992
remote_addr = (host,port)
client = socket.socket()

def init_connection():
    client.connect(remote_addr)
    print('connecting remote socket server')
    print('input key')
    client.send(input('key:>').encode('utf-8'))
    response = client.recv(1024).decode()
    if response.startswith('welcome'):
        print(response)
        return 1
    else:
        return 0

def download(filename):
    if os.path.exists('./download'):
        if client.recv(1024).decode().startswith('filename is not exists'):
            print('check filename')
            return 0
        else:
            file_size = int(client.recv(1024).decode().replace('@',''))
            download_size = 0
            m = hashlib.md5()
            with open('./download/'+filename,'wb') as f:
                chunk = 1024
                while download_size < file_size:
                    if file_size - download_size < 1024:
                        chunk = file_size - download_size
                    data = client.recv(chunk)
                    f.write(data)
                    m.update(data)
                    download_size+=chunk
            file_md5 = client.recv(1024)
            if file_md5 == m.hexdigest():
                print('download susscess')
            print(client.recv(1024).decode())
    else:
        os.system('mkdir ./download')
        download(filename)
                    
def upload(filename):
    if os.path.exists(filename):
        file_size = os.stat(filename).st_size
        client.send(str(file_size).encode('utf-8') + (1024 -len(str(file_size)))*b'@')
        m = hashlib.md5()
        sent_size = 0
        with open(filename,'rb') as f:
            while file_size - sent_size > 1024:
                data = f.read(1024)
            #client.sendall(data)
                m.update(data)
                client.send(data)
                sent_size += 1024
            data = f.read(file_size%1024)
            client.send(data)
            m.update(data)
        time.sleep(0.5)
        client.send(m.hexdigest().encode())
        response = client.recv(1024)
        print(response.decode())


def command():
    print(client.recv(1024).decode())
    while True:
        client.send(input(':>').encode())
        response = client.recv(4096).decode()
        if response.startswith('quit!'):
            break
        print(response)

if init_connection():
    while True:
        print(client.recv(1024).decode().strip(' '))
        func = input(':>')
        client.send(func.encode())
        time.sleep(1)
        func = func.split()
        cmd = func[0]
        cmdargv = func[1:]
        if cmd.startswith('download'):
            download(cmdargv[0])
        elif cmd.startswith('upload'):
            upload(cmdargv[0])
        elif cmd.startswith('cmd'):
            command()
        else:
            print('error input.connect will be done')
            break
        time.sleep(1)
client.close()

服务端

#-*- coding: UTF-8 -*-
import os
import socket
import datetime
import hashlib
import time
#server 
host = '0.0.0.0'
port = 5992
skt = socket.socket()
myaddr = (host,port)
skt.bind(myaddr)
skt.listen()

def download(filenmae):
    if os.path.isfile(filenmae):
        conn.send('begin to download file'.encode())
        with open(filenmae,'rb') as f:
            m = hashlib.md5()
            file_size = os.stat(filenmae).st_size
            sent_size = 0
            conn.send(str(file_size).encode('utf-8') + (1024 -len(str(file_size)))*b'@')
            #begin to send file
            while sent_size < file_size:
                data = f.read(1024)
                if data == 'EOF':
                    break
                conn.send(data)
                m.update(data)
                sent_size += 1024
        time.sleep(0.5)
        conn.send(m.hexdigest().encode())
        conn.send('OK ! '.encode())
    else:
        conn.send('filename is not exists or not file'.encode())

def upload(filename):
    if os.path.exists('./uploads'):
        time_now = datetime.now()
        if '.' in filename:
            file_name,ext = filename.split('.')
            filename_new = file_name + str(time_now.month) + str(time_now.day) + '.' + ext
        else:
            filename_new = filename + str(time_now.month) + str(time_now.day)
        file_size = int(conn.recv(1024).deocode().replace('@',''))
        received_size = 0
        m = hashlib.md5()
        conn.send('begin to upload file!'.encode())
        with open('./uploads/'+filename_new,'wb') as f:
            while received_size < file_size:
                if file_size - received_size > 1024:
                    chunk = 1024
                else:
                    chunk = file_size - received_size
                filedata = conn.recv(chunk)
                received_size += chunk
                f.write(filedata)
                m.update(filedata)
        #conn.send('checking file'.encode())
        file_md5 = conn.recv(32).decode()
        if file_md5 == m.hexdigest():
            conn.send('file upload done! Ur file is named '.encode()+filename_new.encode())
            

    else:
        os.system('mkdir ./uploads')

def command():
    conn.send('R1gelX Welcome!'.encode())
    while True:
        cmd = conn.recv(1024)
        if cmd.startswith('quitcmd'):
            conn.send('quit! exiting command mode'.encode())
            return 0
        result = os.popen(cmd).read()
        if len(result):
            conn.send(result.encode('utf-8'))
        else:
            conn.send('no output command'.encode())


conn,addr = skt.accept()
key = conn.recv(1024).decode()

if key == 'solomanyyds2929':
    conn.send(b'welcome R1gelX socket menu\n')
    while True:
        conn.send('input download|upload|cmd to use'.encode())
        data = conn.recv(1024).decode()
        cmd = data.split()[0]
        cmdargv = data.split()[1:]
        if cmd.startswith('download'):
            download(cmdargv[0])
        elif cmd.startswith('upload'):
            upload(cmdargv[0])
        elif cmd.startswith('cmd'):
            command()
        time.sleep(0.5)

skt.close()

一些问题

在上传文件的时候,利用了sendall()来一次性上传所有的文件,但是在服务器接受的时候,如果文件过大,会导致上传文件不完整。这里我是通过传递的 file_size 来确定文件上传是否完成,服务端每次接受1024直至最后一次接受file_size/1024,理论上来说,确实没有问题,一些小一点的文件确实可以正常的下载上传。还有待解决。