42 练习：利用PyMySQL和socket编写登陆程序

# client.py 客户端
import socket
import struct
import pickle

def my_send(conn, content):
    b_content = pickle.dumps(content)
    b_length = struct.pack('i', len(b_content))
    conn.send(b_length)
    conn.send(b_content)

def my_recv(conn):
    b_length = conn.recv(4)
    length = struct.unpack('i', b_length)[0]
    b_content = conn.recv(length)
    content = pickle.loads(b_content)
    return content

sk = socket.socket()
sk.connect(('127.0.0.1', 9000))

print('1.register')
print('2.login')
print('3.exit')
mode = input('pls select ur mode>>>')
my_send(sk, mode)
if mode == '1':
    username = input('username:')
    password = input('password:')
    user_dic = {'username': username, 'password': password}
    my_send(sk, user_dic)
    if my_recv(sk): print('register success.')
    else: print('register fail.')
elif mode == '2':
    username = input('username:')
    password = input('password:')
    user_dic = {'username': username, 'password': password}
    my_send(sk, user_dic)
    if my_recv(sk):
        print('login success.')
    else:
        print('login fail.')
else: pass

sk.close()


# server.py
import socketserver
import struct
import pickle
import pymysql
import hashlib

KEY = 'This is the key of hashlib.'
global Current_user

def my_send(conn, content):
    b_content = pickle.dumps(content)
    b_length = struct.pack('i', len(b_content))
    conn.send(b_length)
    conn.send(b_content)

def my_recv(conn):
    b_length = conn.recv(4)
    length = struct.unpack('i', b_length)[0]
    b_content = conn.recv(length)
    content = pickle.loads(b_content)
    return content

def login(username, password):
    sql_conn = pymysql.connect(host='127.0.0.1', user='root', password='123', database='user_info')
    cur = sql_conn.cursor()
    md = hashlib.md5(KEY.encode('utf-8'))
    md.update(password.encode('utf-8'))
    md5_pwd = md.hexdigest()
    sql = 'select * from user where username=%s and password=%s'
    cur.execute(sql, (username, md5_pwd))
    if cur.rowcount:
        flag = True
    else: flag = False
    cur.close()
    sql_conn.close()
    return flag

def register(username, password):
    sql_conn = pymysql.connect(host='127.0.0.1', user='root', password='123', database='user_info')
    cur = sql_conn.cursor()
    md = hashlib.md5(KEY.encode('utf-8'))
    md.update(password.encode('utf-8'))
    md5_pwd = md.hexdigest()
    try:
        cur.execute(f'select * from user where username = {username}')
    except Exception: pass
    if cur.rowcount:
        cur.close()
        sql_conn.close()
        return False
    sql = 'insert into user values(%s, %s)'
    try:
        cur.execute(sql, (username, md5_pwd))
        sql_conn.commit()
    except Exception as e: print(e)
    cur.close()
    sql_conn.close()
    return True

class Myserver(socketserver.BaseRequestHandler):
    def handle(self):
        conn = self.request  # self.request 就是 conn
        mode = my_recv(conn)
        if mode == '1':
            user_dic = my_recv(conn)
            ret = register(user_dic['username'], user_dic['password'])
            my_send(conn, ret)
        elif mode == '2':
            user_dic = my_recv(conn)
            ret = login(user_dic['username'], user_dic['password'])
            if ret:
                global Current_user
                Current_user = user_dic['username']
            my_send(conn, ret)
        conn.close()



server = socketserver.ThreadingTCPServer(('127.0.0.1', 9000), Myserver)
server.serve_forever()