一个基于python的即时通信程序

广播信息（上线）:
{
  'status': 信息状态标志,
  'user_info': 本机的用户名和主机名,
  'pub_key': 本机生成的公钥,
}
广播信息（下线）：
{
  'status': 信息状态标志,
  'user_info': 本机的用户名和主机名,
}

用户列表的元素：
{
  'user_info': 对应用户的用户名和主机名,
  'pub_key': 对应用户的公钥,
  'addr': 用户对应的ip,
}

信息确认列表的元素：
{
  'confirm_seq': 信息序列号,
  'user': 发送信息的用户的用户名,
  'msg': 发送的信息,
  'addr': 信息的目的ip和端口,
}


通信信息：
{
  'status': 信息序列号,
  'user': 发送信息的用户的用户名,
  'msg': 发送的信息,
}

!/usr/bin/env python
#coding=utf-8
#author: cjyfff
#blog: http://www.cnblogs.com/cjyfff/

import socket
import os
import threading
import traceback
import rsa

user_list = []
confirm_list = []
username = os.environ['USER']
hostname = os.popen('hostname').read()
(pubkey, privkey) = rsa.newkeys(1024)
pub = pubkey.save_pkcs1()


class MyThread(threading.Thread):
    '''这个类用于创建新的线程'''

    def __init__(self, func, args, name=''):
        threading.Thread.__init__(self)
        self.name = name
        self.func = func
        self.args = args

    def run(self):
        apply(self.func, self.args)


def broadcast(broADDR, status):
    '''发送广播信息模块
    用于发送广播信息给其他主机，通知其他主机本主机上线下线状态，以及发送本机的信息给其他主机。
    这个模块会在广播信息前添加上status这个参数的值。在本程序中，当需要通知其他主机，本机已经上线时，
    会传递"online"给status，当需要通知其他主机本机即将下线时，会传递"offline"给status。
    '''
    global username, hostname, pub

    def broadcast_send(oMsg):
        udpSock2 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udpSock2.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        udpSock2.sendto(oMsg, broADDR)
    oMsg = {}
    if status == 'online':
        oMsg = {
            'status': status,
            'user_info': ' '.join([username, hostname]),
            'pub_key': pub,
        }
        broadcast_send(str(oMsg))
    elif status == 'offline':
        oMsg = {
            'status': status,
            'user_info': ' '.join([username, hostname]),
        }
        broadcast_send(str(oMsg))


def recv_msg(localADDR, BUFSIZ, udpSock, port):
    '''信息接收模块
    这个模块的主要功能是，跟据接收到的广播信息更新用户列表，以及处理对端发送过来信息
    '''
    global user_list, confirm_list, username, hostname, pub, privkey

    while True:
        try:
            data, addr = udpSock.recvfrom(BUFSIZ)
        except:
            break
        if data:
            data = eval(data)
        addr = addr[0]
        if data['status'] == 'online':
            user_list_info = {
                'user_info': data['user_info'],
                'pub_key': data['pub_key'],
                'addr': addr,
            }
            if user_list_info not in user_list:
                user_list.append(user_list_info)
                # 把对方添加上用户列表的同时，还要把自己的信息发给对方，以便对方更新用户列表
                respond_msg = {
                    'status': 'respon_online',
                    'user_info': ' '.join([username, hostname]),
                    'pub_key': pub,
                }
                udpSock.sendto(str(respond_msg), (addr, port))

        elif data['status'] == 'offline':
            user_list_info = {
                'user_info': data['user_info'],
            }
            for i in xrange(len(user_list)):
                for k, v in user_list[i].iteritems():
                    if user_list_info['user_info'] == v:
                        del user_list[i]

        elif data['status'] == 'respon_online':
            user_list_info = {
                'user_info': data['user_info'],
                'pub_key': data['pub_key'],
                'addr': addr,
            }
            if user_list_info not in user_list:
                user_list.append(user_list_info)

        elif data['status'] == 'quit':
            print "对方已断开连接，请输入'quit'或'q'返回主菜单"
            continue

        elif data['status'] == 'local_quit':
            continue

        else:
            confirm_msg = data['status']
            # 假如收到的确认标志和确认表中的某项匹配，删除该项
            for i in xrange(len(confirm_list)):
                if confirm_list[i]['confirm_seq'] == confirm_msg:
                    del confirm_list[i]
            if not data['msg']:
                continue
            addr_list = []
            for x in user_list:
                # 提取出用户表中所有用户的地址，存到addr_list中：
                addr_list.append(x['addr'])

            # 检查发送信息的用户的地址是否在用户列表当中：
            if addr in addr_list:
                # 反馈收到确认信息给对方：
                confirm_res = {'status': confirm_msg, 'msg': 0}
                udpSock.sendto(str(confirm_res), (addr, port))
                # 打印信息：
                data_user = data['user']
                try:
                    data_msg = rsa.decrypt((data['msg']), privkey)
                except DecryptionError:
                    print "解码出现异常，请重新连接"
                    continue
                print data_user, ":", data_msg


def print_userlist():
    '''打印用户列表模块'''
    global user_list
    user_list_len = len(user_list)
    print "当前有%d个用户在线：" % user_list_len
    for i in xrange(user_list_len):
        print "ID:", i+1, ":", user_list[i]['user_info'].strip('
'), 
            "come from:", user_list[i]['addr']


def send_msg(udpSock, cli_addr, cli_pub_key, port):
    '''信息发送模块'''
    import random
    global username, confirm_list
    quit_list = ['q', 'quit', 'exit']
    cli_pub_key_rip = rsa.PublicKey.load_pkcs1(cli_pub_key)

    while True:
        msg = raw_input("> ")
        if msg in quit_list:
            # quit_msg_to_local用于通知本机对话结束，回收socket
            quit_msg_to_local = {'status': 'local_quit'}
            quit_msg_to_cli = {'status': 'quit'}
            udpSock.sendto(str(quit_msg_to_local), ('localhost', port))
            udpSock.sendto(str(quit_msg_to_cli), cli_addr)
            break

        random_num = random.randint(0, 1000)
        msg = rsa.encrypt(msg, cli_pub_key_rip)
        output_msg = {
            'status': str(random_num),
            'user': username,
            'msg': msg,
        }
        confirm_list_member = {
            'confirm_seq': str(random_num),
            'user': username,
            'msg': msg,
            'addr': cli_addr,
        }
        confirm_list.append(confirm_list_member)

        udpSock.sendto(str(output_msg), cli_addr)


def confirm_successd(udpSock):
    '''确认信息到达模块
    采用类似于最久未使用（LRU）算法，每隔5秒钟检查一下信息确认列表(confirm_list)，当信息确认列表长度大于5时（
    也就是说未确认接收的信息大于5），把信息确认列表前一半的信息再一次发送。
    '''
    import time
    global confirm_list

    while True:
        confirm_list_len = len(confirm_list)
        if confirm_list_len > 5:
            for i in xrange(confirm_list_len/2):
                repeat_output_msg = {
                    'status': confirm_list[i]['confirm_seq'],
                    'user': confirm_list[i]['user'],
                    'msg': confirm_list[i]['msg'],
                }
                #msg = confirm_list[i][0]
                addr = confirm_list[i]['addr']
                udpSock.sendto(str(repeat_output_msg), addr)
            time.sleep(5)
        else:
            time.sleep(5)


def option(udpSock, BUFSIZ, broADDR, port):
    '''选项菜单模块'''
    while True:
        print '''
        请输入您的选项:
        1 显示用户列表
        2 连接到指定用户，并开始对话
        3 退出
        '''
        action = raw_input("> ")
        if action is '1':
            print_userlist()

        elif action is '2':
            client_option = raw_input("您想连接到哪个用户？，请输入对应的id号：
")
            try:
                # 获取对端的地址
                cli_addr = (user_list[int(client_option)-1]['addr'], port)
                cli_pub_key = user_list[int(client_option)-1]['pub_key']
            except IndexError:
                print "没有这个用户，请重新选择："
                continue
            print "已建立好连接，可以开始对话，输入quit或q可以结束会话"
            threads = []
            t2 = MyThread(send_msg, (udpSock, cli_addr, cli_pub_key, port), send_msg.__name__)
            threads.append(t2)
            t3 = MyThread(confirm_successd, (udpSock, ), confirm_successd.__name__)
            threads.append(t3)
            for t in threads:
                t.setDaemon(True)
                t.start()
            t2.join()#send_msg中止之前，让父线程一直在阻塞状态
            print "连接中断，返回主菜单"

        elif action is '3':
            broadcast(broADDR, 'offline')
            udpSock.close()
            print "再见！"
            break

        else:
            pass


def main():
    '''主函数'''
    host = ''
    port = 2425
    broADDR = ('<broadcast>', port)
    localADDR = (host, port)
    BUFSIZ = 1024
    try:
        broadcast(broADDR, 'online')
        udpSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udpSock.bind(localADDR)
        t1 = MyThread(recv_msg, (localADDR, BUFSIZ, udpSock, port, ),
                       recv_msg.__name__)
        t1.setDaemon(True)
        t1.start()
        option(udpSock, BUFSIZ, broADDR, port)
    except (KeyboardInterrupt, SystemError):
        udpSock.close()
        raise
    except:
        traceback.print_exc

if __name__ == '__main__':
    main()