Python 作业#4

作业题目

在作业#3的基础上实现localProxy命令行参数账号登录

在作业#3的基础上实现remoteProxy多账号认证

remoteProxy采用SQLite3数据库进行用户账号管理（用户名、密码）

remoteProxy使用aiosqlite操作SQLite3数据库

# dualProxy.py
import asyncio
import argparse
import sys
import logging
import ipaddress
import msvcrt
import aiosqlite
from struct import unpack, pack

pass_words = ""

def pwd_input():
    chars = []
    while True:
        try:
            newChar = msvcrt.getch().decode(encoding="utf-8")
        except:
            return input()
        if newChar in '\r\n':
             break
        elif newChar == '\b':
             if chars:
                 del chars[-1]
                 msvcrt.putch('\b'.encode(encoding='utf-8'))
                 msvcrt.putch( ' '.encode(encoding='utf-8'))
                 msvcrt.putch('\b'.encode(encoding='utf-8'))
        else:
            chars.append(newChar)
            msvcrt.putch('*'.encode(encoding='utf-8'))
    return (''.join(chars) )


async def io_close(writer):
    if not writer:
        await asyncio.sleep(0.0001)
        return

    host, port = writer.get_extra_info("peername")
    log.info(f"close {host} {port}")
    try:
        writer.close()
        await writer.wait_close()
    except Exception as exc:
        pass


async def msg_transfer(reader, writer):
    while reader.at_eof:
        try:
            data = await reader.read(1024 * 64)
            if not data:
                await io_close(writer)
                break
        except (ConnectionAbortedError, ConnectionResetError) as e:
            log.info(f"exit exceptionally {repr(e)}")
            await io_close(writer)
            break

        try:
            writer.write(data)
            await writer.drain()
        except (ConnectionAbortedError, ConnectionResetError) as e:
            log.info(f"exit exceptionally {repr(e)}")
            await io_close(writer)
            break

async def verify(user_name, pass_word, writer):

    async with aiosqlite.connect('password.db') as db:
        async with db.execute(f"SELECT pass_word FROM information WHERE user_name = '{user_name}'") as  cursor:
            async for row in cursor:

                if row[0] == pass_word:
                    flag = True
                else:
                    flag = False

    if flag == False:
        data = "wrong pass word"
        log.info(f"{data}")
        writer.write(data.encode())
        await writer.drain()
        await io_close(writer)
    else:
        data = "verification passed"
        log.info(f"{data}")
        writer.write(data.encode())
        await writer.drain()

    return flag


async def handle_local(reader, writer):
    global pass_words

    data = await reader.read(1024 * 64)

    if data[0] == 5:
        proxy_type = "socks5"
    else:
        proxy_type = "https"


    if (proxy_type == "socks5"):
        host, port = writer.get_extra_info('peername')
        log.info(f"connect from {host} {port}")

        writer.write(b'\x05\x00')
        await writer.drain()

        data = await reader.read(1024 * 64)

    else:
        host, port = writer.get_extra_info('peername')
        log.info(f"connect from {host} {port}")


    try:
        reader_remote, writer_remote = await asyncio.open_connection(args.remote_host, args.remote_port)
        log.info(f"connect to {args.remote_host} {args.remote_port} successfully")
    except:
        log.info(f"connect to {args.remote_host} {args.remote_port} unsuccessfully")
        await io_close(writer)
        return

    info = data + (args.user_name + ',' + pass_words).encode()
    writer_remote.write(info)
    await writer.drain()

    response = await reader_remote.read(1024 * 64)
    log.info(f"{response.decode()}")
    if response.decode() == "wrong pass word":
        print('retype password:', end=' ', flush=True)
        pass_words = pwd_input()
        print('')

        info = data + (args.user_name + ',' + pass_words).encode()
        return

    client_to_remote = msg_transfer(reader, writer_remote)
    remote_to_client = msg_transfer(reader_remote, writer)
    await asyncio.gather(client_to_remote, remote_to_client)


async def handle_remote(reader, writer):
    data = await reader.read(1024 * 64)

    if data[0] == 5:
        proxy_type = "socks5"
    else:
        proxy_type = "https"

    reader_remote = None
    writer_remote = None

    user_name = None
    pass_word = None

    if(proxy_type == "socks5"):
        host, port = writer.get_extra_info('peername')
        log.info(f"connect from {host} {port}")

        info = unpack('!4B', data[:4])
        ver, cmd, atyp = info[0], info[1], info[3]


        # domain name
        if ver == 5 and cmd == 1 and atyp == 3:

            addr_len = unpack('!B', data[4:5])[0]
            dst_addr = data[5:addr_len + 5].decode()
            dst_port = unpack('!H', data[addr_len + 5: addr_len + 7])[0]
            log.info(f'destination domain name {dst_addr}，destination port {dst_port}')

            msg = data[addr_len + 7:].decode()
            user_name = msg.split(',')[0]
            pass_word = msg.split(',')[1]

            if await verify(user_name, pass_word, writer) == False:
                return

            try:
                reader_remote, writer_remote = await asyncio.open_connection(dst_addr, dst_port)
                writer.write(pack('!5B', 5, 0, 0, 3, addr_len) + dst_addr.encode() + pack('!H', dst_port))
                await writer.drain()
                log.info(f'connect to {dst_addr} {dst_port} successfully')
            except (TimeoutError, ConnectionRefusedError, OSError) as e:
                log.info(f'connect to {dst_addr} {dst_port} unsuccessfully')
                writer.write(pack('!5B', 5, 3, 0, 3, addr_len) + dst_addr.encode() + pack('!H', dst_port))
                await writer.drain()
                await io_close(writer)
                return


        # ipv4
        elif ver == 5 and cmd == 1 and atyp == 1:

            dst_address = str(ipaddress.ip_address(data[4:8]))
            dst_port = unpack('H', data[8:10])[0]
            log.info(f'destination address {dst_address}，destination port {dst_port}')

            msg = data[10:].decode()
            user_name = msg.split(',')[0]
            pass_word = msg.split(',')[1]

            if await verify(user_name, pass_word, writer) == False:
                return

            try:
                reader_remote, writer_remote = await asyncio.open_connection(dst_address, dst_port)
                writer.write(pack('!8B', 5, 0, 0, 1, *unpack('!BBBB', data[4:8])) + pack('!H', dst_port))
                await writer.drain()
                log.info(f'connect to {dst_addr} {dst_port} successfully')
            except (TimeoutError, ConnectionRefusedError, OSError) as e:
                log.info(f'connect to {dst_addr} {dst_port} unsuccessfully')
                writer.write(pack('!8B', 5, 3, 0, 1, *unpack('!BBBB', data[4:8])) + pack('!H', dst_port))
                await writer.drain()
                await io_close(writer)
                return


        # ipv6
        elif ver == 5 and cmd == 1 and atyp == 4:
            dst_address = str(ipaddress.ip_address(data[4:20]))
            dst_port = unpack('H', data[20:22])[0]
            log.info(f'destination address {dst_address}，destination port {dst_port}')

            msg = data[22:].decode()
            user_name = msg.split(',')[0]
            pass_word = msg.split(',')[1]

            if await verify(user_name, pass_word, writer) == False:
                return

            try:
                reader_remote, writer_remote = await asyncio.open_connection(dst_address, dst_port)
                writer.write(pack('!20B', 5, 0, 0, 1, *unpack('!16B', data[4:20])) + pack('!H', dst_port))
                await writer.drain()
                log.info(f'connect to {dst_addr} {dst_port} successfully')
            except (TimeoutError, ConnectionRefusedError, OSError) as e:
                log.info(f'connect to {dst_addr} {dst_port} unsuccessfully')
                writer.write(pack('!20B', 5, 3, 0, 1, *unpack('!16B', data[4:20])) + pack('!H', dst_port))
                await writer.drain()
                await io_close(writer)
                return


    else:
        host, port = writer.get_extra_info('peername')
        log.info(f"connect from {host} {port}")

        info = data.decode().split("\r\n")
        request = info[0].split(" ")

        method = request[0]
        host = request[1].split(":")[0]
        port = request[1].split(":")[1]
        ver = request[2]
        log.info(f'destination domain name {host}，destination port {port}')

        msg = info[-1]
        user_name = msg.split(',')[0]
        pass_word = msg.split(',')[1]

        if await verify(user_name, pass_word, writer) == False:
            return

        if method.upper() == 'CONNECT':

            try:
                reader_remote, writer_remote = await asyncio.open_connection(host, port)
                reply = ver + " 200 Connection established\r\n"
                reply += "Proxy-agent: wyx\r\n\r\n"
                writer.write(reply.encode())
                await writer.drain()
                log.info(f'connect to {host} {port} successfully')
            except (TimeoutError, ConnectionRefusedError) as e:
                log.info(f'connect to {host} {port} unsuccessfully')
                await io_close(writer)
                return


    client_to_remote = msg_transfer(reader, writer_remote)
    remote_to_client = msg_transfer(reader_remote, writer)
    await asyncio.gather(client_to_remote, remote_to_client)


async def local():
    print('password:', end=' ', flush=True)
    global pass_words
    pass_words = pwd_input()
    print('')

    server = await asyncio.start_server(handle_local, '127.0.0.1', args.listen_port)

    host, port = server.sockets[0].getsockname()
    log.info(f'Serving on {host} {port}')

    async with server:
        await server.serve_forever()

async def remote():

    server = await asyncio.start_server(handle_remote, args.listen_host, args.listen_port)

    host, port = server.sockets[0].getsockname()
    log.info(f'Serving on {host} {port}')

    async with server:
        await server.serve_forever()


async def main():
    if args.remote_host:
        asyncio.create_task(local())
    else:
        asyncio.create_task(remote())

    while True:
        await asyncio.sleep(1)


if __name__ == "__main__":
    _logFmt = logging.Formatter('%(asctime)s %(levelname).1s %(lineno)-3d %(funcName)-20s %(message)s', datefmt='%H:%M:%S')
    _consoleHandler = logging.StreamHandler()
    _consoleHandler.setLevel(logging.DEBUG)
    _consoleHandler.setFormatter(_logFmt)

    log = logging.getLogger(__file__)
    log.addHandler(_consoleHandler)
    log.setLevel(logging.DEBUG)

    _parser = argparse.ArgumentParser(description='local socks5 https dual proxy server')

    _parser.add_argument('--host', dest='listen_host', metavar='listen_host', default='127.0.0.1', help='proxy listen host default 127.0.0.1')
    _parser.add_argument('--port', dest='listen_port', metavar='listen_port', required=True, help='proxy listen port')
    _parser.add_argument('--rhost', dest='remote_host', metavar='remote_host', nargs='?', default=None, help='remote proxy host')
    _parser.add_argument('--rport', dest='remote_port', metavar='remote_port', nargs='?', default=None, help='remote proxy port')
    _parser.add_argument('--user', dest='user_name', metavar='user_name', nargs='?', default=None, help='user name')

    args = _parser.parse_args()

    if args.remote_host:
        if not args.remote_port or not args.user_name:
            _parser.print_help()
            exit(1)

    asyncio.run(main())

    
#sql.py
import aiosqlite
import asyncio

async def main():
    async with aiosqlite.connect('password.db') as db:
        await db.execute("CREATE TABLE information(user_name VARCHAR(10), pass_word VARCHAR(10))")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u1', 'wWOjNqRM')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u2', 'IUcLlsbA')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u3', 'UQ9UKQ51')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u4', 'los8PDQA')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u5', '5m5SYYp5')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u6', 'EFciaMwB')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u7', 'eBZCRUox')")
        await db.execute("INSERT INTO information(user_name, pass_word) VALUES('u8', '5ZG2AJgc')")
        await db.commit()

        async with db.execute("SELECT user_name, pass_word FROM information") as cursor:
            async for row in cursor:
                print(f'f1={row[0]} f2={row[1]}')

async def main_logic():
    asyncio.create_task(main())

    while True:
        await asyncio.sleep(1)

if __name__ == '__main__':
    asyncio.run(main_logic())