1. 程式人生 > 程式設計 >基於socket的簡單聊天程式

基於socket的簡單聊天程式


title: 基於socket的簡單聊天程式 date: 2019-04-10 16:55:14 tags:

  • Network

這是一個可以自由切換聊天物件,也能傳輸檔案,還有偽雲端聊天資訊暫存功能的聊天程式!( ̄▽ ̄)"


去年暑假的小學期的電子設計課上我用STC與電腦相互通訊,製作出了一個Rader專案(該專案的完整程式碼在我的GitHub上)。這個專案地大致思想是:下位機(STC)使用步進電機帶動超聲波模組採集四周的距離,然後用485序列匯流排上傳到上位機(電腦),上位機將這些資料收集並繪製略醜的雷達圖。由於上下位機處理資料的速度不一致,容易導致不同步的現象。當時為瞭解決這個問題,用了一個簡單的方法,現在發現這個方法和“停等協議”十分相似。

這學期的計網實驗要求基於socket傳輸資料,相比於在485匯流排上實現停等協議,socket還是很簡單的。

Naive版聊天程式

最簡單的socket通訊程式只需要兩個程式就可以跑起來了,一個作為服務端,另一個作為客戶端,然後兩者之間傳輸資料。

# Server
import socket
from socket import AF_INET,SOCK_STREAM

serverSocket = socket.socket(AF_INET,SOCK_STREAM)
srv_addr = ("127.0.0.1",8888)
serverSocket.bind(srv_addr)
serverSocket.listen()

print("[Server INFO] listening..."
) while True: conn,cli_addr = serverSocket.accept() print("[Server INFO] connection from {}".format(cli_addr)) message = conn.recv(1024) conn.send(message.upper()) conn.close() 複製程式碼
# Client
import socket
from socket import AF_INET,SOCK_STREAM

clientSocket = socket.socket(AF_INET,8888
) print("[Client INFO] connect to {}".format(srv_addr)) clientSocket.connect(srv_addr) message = bytes(input("Input lowercase message> "),encoding="utf-8") clientSocket.send(message) modifiedMessage = clientSocket.recv(1024).decode("utf-8") print("[Client INFO] recv: '{}'".format(modifiedMessage)) clientSocket.close() 複製程式碼

多使用者版

上面這種模式是十分naiive的。比如為了切換使用者(假設不同使用者在不同的程式上),就只能先kill原先的程式,然後修改程式碼中的IP和Port,最後花了1分鐘時間才能開始聊天。而且這種方式最大的缺陷是隻有知道了對方的IP和Port之後才能開始聊天。

為瞭解決Naiive版聊天程式的缺點,可以構建如下C/S拓撲結構。

這個拓撲的結構的核心在於中央的Server,所有Client的連線資訊都會被儲存在Server上,Server負責將某個Client的聊天資訊轉發給目標Client。

資料格式設計

像TCP協議需要報文一樣,這個簡單聊天程式的資訊轉發也需要Server識別每一條資訊的目的,才能準確轉發資訊。這就需要設計協議報文的結構(顯然這是在應用層上的實現)。由於應用場景簡單,我是用的協議結構如下:

sender|receiver|timestamp|msg
複製程式碼

這是一個四元組,每個元素用管道符|分割。具體來說每個Client(客戶程式)傳送資料給Server之前都會在msg之前附加:傳送方標識sender、接受方標識receiver以及本地時間戳timestamp。對應的程式碼端如下:

info = "{}|{}|{}|{}".format(self.UserID,targetID,timestamp,msg)
複製程式碼

這樣Server接收到報文之後就能“正確”轉發訊息了。

這裡的“正確”被加上了引號,這是為什麼?因為在我設計該乞丐版協議的時候簡化場景中只存在唯一使用者ID的場景,如果有個叫“Randool”的使用者正在和其他使用者聊天,這個時候另一個“Randool”登陸了聊天程式,那麼前者將不能接收資訊(除非再次登入)。不過簡單場景下還是可以使用的。

解決方法可以是在Client登入Server時新增驗證的步驟,讓重複使用者名稱無法通過驗證。

訊息佇列

該聊天程式使用的傳輸層協議是TCP,這是可靠的傳輸協議,但聊天程式並不能保證雙方一定線上吧,聊天一方在任何時候都可以退出聊天。但是一個健壯的聊天程式不能讓資訊有所丟失,由於傳輸層已經不能確保資訊一定送達,那麼只能寄希望於應用層。

由於訊息是通過Server轉發的,那麼只要在Server上為每一個Client維護一個訊息佇列即可。資料結構如下:

MsgQ = {}
Q = MsgQ[UserID]
複製程式碼

使用這種資料結構就可以模擬雲端聊天記錄暫存的功能了!

檔案傳輸

檔案傳輸本質上就是傳輸訊息,只不過檔案傳輸的內容不是直接顯示在螢幕上罷了。相比於純聊天記錄的傳輸,檔案傳輸需要多附加上檔名,

base64編碼傳輸

普通的聊天資訊中不會出現管道符,但是程式碼和字元表情就不一定了∑( 口 ||,如果資訊中出現了管道符就會導致協議解析失效,因此需要一種方法將msg中的|隱藏掉。思路是轉義,但是這個需要手工重寫協議解析程式碼,不夠美觀。由於之前瞭解過資訊保安中的相關知識,還記得有一種編碼方式是base64,由於base64編碼結果不會出現管道符,那麼問題就簡單了,只需要用base64將傳輸資訊重新編碼一番。並且這是一種“即插即用”的方式,只要自定義base64的編碼解碼函式,然後巢狀在待傳送msg的外面即可。

import base64

b64decode = lambda x: base64.b64decode(x.encode()).decode()
b64encode = lambda x: base64.b64encode(x.encode()).decode()
複製程式碼

將傳送資訊改寫為如下形式:

info = "{}|{}|{}|{}||".format(self.UserID,b64encode(msg))
複製程式碼

終端高亮顯示

樸素的文字列印在螢幕上難以區分主次,使用者體驗極差,因此可以使用終端高亮的方法凸顯重要資訊。在網上查到了一種高亮的方式,但是僅限於Linux系統。其高亮顯示的格式如下:

\033[顯示方式;前景色;背景色mXXXXXXXX\033[0m

中間的XXXXXXXX就是需要顯示的文字部分了。顯示方式,前景色,背景色是可選引數,可以只寫其中的某一個;另外由於表示三個引數不同含義的數值都是唯一的沒有重複的,所以三個引數的書寫先後順序沒有固定要求,系統都能識別;但是,建議按照預設的格式規範書寫。

這個部分參考了Python學習-終端字型高亮顯示,因此對於引數的配置方面不再多說

效果

有多種終端分屏外掛,這裡推薦tmux,上面的分屏效果使用的就是tmux

程式碼實現

服務端程式碼

import queue
import socket
import time

import _thread


hostname = socket.gethostname()
port = 12345

"""
The info stored in the queue should be like this:
"sender|receiver|timestamp|msg"
and all item is str.
"""
MsgQ = {}


def Sender(sock,UserID):
    """
    Fetch 'info' from queue send to UserID.
    """
    Q = MsgQ[UserID]
    try:
        while True:
            # get methord will be blocked if empty
            info = Q.get()
            sock.send(info.encode())
    except Exception as e:
        print(e)
        sock.close()
        _thread.exit_thread()


def Receiver(sock):
    """
    Receive 'msg' from UserID and store 'info' into queue.
    """
    try:
        while True:
            info = sock.recv(1024).decode()
            print(info)
            info_unpack = info.split("|")
            receiver = info_unpack[1]
            
            exit_cmd = receiver == "SEVER" and info_unpack[3] == "EXIT"
            assert not exit_cmd,"{} exit".format(info_unpack[0]) 
            
            if receiver not in MsgQ:
                MsgQ[receiver] = queue.Queue()
            MsgQ[receiver].put(info)

    except Exception as e:
        print(e)
        sock.close() 
        _thread.exit_thread()


class Server:
    def __init__(self):
        self.Sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        self.Sock.bind((hostname,port))
        self.Sock.listen()
        # self.threads = []

    def run(self):
        print("\033[35;40m[ Server is running ]\033[0m")
        # print("[ Server is running ]")
        while True:
            sock,_ = self.Sock.accept()

            # Register for new Client
            UserID = sock.recv(1024).decode()
            print("Connect to {}".format(UserID))

            # Build a message queue for new Client
            if UserID not in MsgQ:
                MsgQ[UserID] = queue.Queue()

            # Start two threads
            _thread.start_new_thread(Sender,(sock,UserID))
            _thread.start_new_thread(Receiver,))

    def close(self):
        self.Sock.close()


if __name__ == "__main__":
    server = Server()
    try:
        server.run()
    except KeyboardInterrupt as e:
        server.close()
        print("Server exited")
複製程式碼

客戶端程式碼

import socket
import sys,os
import time
import base64
import _thread
from SktSrv import hostname,port

b64decode = lambda x: base64.b64decode(x.encode()).decode()
b64encode = lambda x: base64.b64encode(x.encode()).decode()


def Receiver(sock):
    from_id = ""
    fr = None   # file handle
    while True:
        info = sock.recv(1024).decode()
        info_unpacks = info.split("||")[:-1]
        for info_unpack in info_unpacks:
            sender,_,msg = info_unpack.split("|")
            msg = b64decode(msg)    # base64解碼

            # Start a new session
            if from_id != sender:
                from_id = sender
                print("==== {} ====".format(sender))
            
            if msg[:5] == "@FILE":  # FILENAME,FILE,FILEEND
                # print(msg)
                if msg[:10] == "@FILENAME:":
                    print("++Recvive {}".format(msg[9:]))
                    fr = open(msg[10:]+".txt","w")
                elif msg[:9] == "@FILEEND:":
                    fr.close()
                    print("++Recvive finish")
                elif msg[:6] == "@FILE:":
                    fr.write(msg[6:])
                continue

            show = "{}\t{}".format(timestamp,msg)
            print("\033[1;36;40m{}\033[0m".format(show))


class Client:
    def __init__(self,UserID: str=None):
        if UserID is not None:
            self.UserID = UserID
        else:
            self.UserID = input("login with userID >> ")
        self.Sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        self.server_addr = (hostname,port)

    def Sender(self):
        """ 
        Send info: "sender|receiver|timestamp|msg"
        Change to name: '@switch:name'
        Trans file: '@trans:filename'
        """
        targetID = input("Chat with > ")
        while True:
            msg = input()
            if not len(msg):
                continue

            lt = time.localtime()
            timestamp = "{}:{}:{}".format(lt.tm_hour,lt.tm_min,lt.tm_sec)

            if msg == "@exit":          # 退出
                print("Bye~")
                return

            elif msg == "@help":
                continue

            elif msg[:8] == "@switch:": # 切換聊天物件
                targetID = msg.split(":")[1]
                print("++Switch to {}".format(targetID))
                continue
            
            elif msg[:7] == "@trans:":  # 傳送檔案
                filename = msg.split(":")[1]
                if not os.path.exists(filename):
                    print("!!{} no found".format(filename))
                    continue
                print("++Transfer {} to {}".format(filename,targetID))
                head = "{}|{}|{}|{}||".format(self.UserID,b64encode("@FILENAME:"+filename))
                self.Sock.send(head.encode())
                with open(filename,"r") as fp:
                    while True:
                        chunk = fp.read(512)
                        if not chunk:
                            break
                        chunk = "{}|{}|{}|{}||".format(self.UserID,b64encode("@FILE:"+chunk))
                        self.Sock.send(chunk.encode())
                tail = "{}|{}|{}|{}||".format(self.UserID,b64encode("@FILEEND:"+filename))
                self.Sock.send(tail.encode())
                print("++Done.")
                continue
            
            info = "{}|{}|{}|{}||".format(self.UserID,b64encode(msg))
            self.Sock.send(info.encode())

    def run(self):
        try:
            self.Sock.connect(self.server_addr)
            print("\033[35;40m[ Client is running ]\033[0m")
            # print("[ Client is running ]")

            # Register UserID
            self.Sock.send(self.UserID.encode())

            # Start Receiver threads
            _thread.start_new_thread(Receiver,(self.Sock,))
            self.Sender()   # Use for Send message

        except BrokenPipeError:
            print("\033[1;31;40mMissing connection\033[0m")

        finally:
            print("\033[1;33;40mYou are offline.\033[0m")
            self.exit_client()
            self.Sock.close()

    def exit_client(self):
        bye = "{}|{}|{}|{}".format(self.UserID,"SEVER","","EXIT")
        self.Sock.send(bye.encode())


if __name__ == "__main__":
    client = Client()
    client.run()
複製程式碼

P2P版

上面的多使用者版聊天程式雖然可以實現靈活的使用者切換聊天功能,但是實際上由於所有的資料都會以伺服器為中轉站,會對伺服器造成較大的壓力。更加靈活的結構是使用P2P的方式,資料只在Client間傳輸。應該是將伺服器視為類似DNS伺服器的角色,只維護一個Name <--> (IP,Port)的查詢表,而將連線資訊轉移到Client上。

存在的問題

P2P版本的聊天程式並不只是實現上述的功能就可以了,考慮到前邊“訊息佇列”中實現的功能:在使用者退出後,聊天資訊需要能儲存在一個可靠的地方。既然聊天雙方都存在退出的可能,那麼在這個場景下這個“可靠的地方”就是伺服器了。這也就是說P2P版本的Client除了建立與其他Client之間的TCP連線,還需要一直保持和Server的連線!

注意這一點,之前是為了減輕Server的壓力,減少連線的數量才使用P2P的模式的,但是在該模式為了實現“訊息佇列”的功能卻還是需要Server儲存連線。

改進方式

如果要進一步改善,可以按照下面的方式:

  1. Client C1登入時與Server建立連線,Server驗證其登入合法性,然後斷開連線。
  2. C1選擇聊天物件C2,C2的IP等資訊需要從Server中獲取,因此C1再次建立與Server的連線,完成資訊獲取後,斷開連線。
  3. C1與C2的正常聊天資訊不通過Server,而是真正的P2P傳輸。
  4. 聊天一方意外斷開後(假設為C2),C1傳輸的資訊無法到達,並且C1可以感知到資訊無法到達;這個時候C1再次建立與Server的連線,將未能送達的資訊儲存到Server上的“訊息佇列”。

補充一點:在步驟2中,如果C2未上線或C2意外斷開,由於Server並不能及時知道Client的資訊,因此需要“心跳包機制”,Client登入後定時向Server傳送alive資訊,Server收到資訊後維持或更新資訊。

這樣Server從始至終沒有一直維持著連線,連線數量是動態變化的,在查詢併發量較小的情況下對伺服器資源的利用率是很小的。

進一步可以思考什麼?

如果有多個Server,如何規劃Server之間的拓撲?比如Fat-Tree之類的...