Python黑客编程之类nc工具

目的

用python socket编写一款类似NetCat的工具，可以在服务器上远程执行命令，从服务器上下载文件

代码

服务端和客户端用同一套代码，用-l参数进行区分

import argparse
import shlex
import socket
import subprocess
import sys
import textwrap
import threading


class NetCat:
    def __init__(self, args, buffer=None):
        self.args = args
        self.buffer = buffer
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    def run(self):
        if self.args.listen:
            self.listen()
        else:
            self.send()

    def send(self):
        self.socket.connect((self.args.target, self.args.port))
        if self.buffer:
            self.socket.send(self.buffer)

        if self.args.file:
            file_buffer = b''
            while True:
                data = self.socket.recv(4096)
                data_len = len(data)
                file_buffer += data
                if data_len < 4096:
                    break
            with open(self.args.file, 'wb') as fp:
                fp.write(file_buffer)
            message = "Successfully send file {} to {}".format(self.args.file, socket.gethostbyname(socket.gethostname()))
            self.socket.send(message.encode())
            print("Successfully received file {}".format(self.args.file))

        if not self.args.command:
            return

        try:
            while True:
                print("<NetCat: #>", end="")
                buffer = input()
                buffer += '\n'
                self.socket.send(buffer.encode())
                data_len = 1
                content = ''
                while data_len:
                    data = self.socket.recv(4096)
                    data_len = len(data)
                    content += data.decode()
                    if data_len < 4096:
                        break
                if content:
                    print(content, end="")

        except KeyboardInterrupt:
            print('User terminated.')
            self.socket.close()
            sys.exit()

    def listen(self):
        self.socket.bind((self.args.target, self.args.port))
        self.socket.listen(5)

        while True:
            client_socket, _ = self.socket.accept()
            client_thread = threading.Thread(target=self.handle_client, args=(client_socket,))
            client_thread.start()

    def handle_client(self, client_socket):
        if self.args.file:
            with open(self.args.file, 'rb') as fp:
                client_socket.send(fp.read())
            recv_buffer = client_socket.recv(64)
            print(recv_buffer.decode())

        if self.args.execute:
            result = self.execute(self.args.execute)
            client_socket.send(result.encode())

        if not self.args.command:
            return

        if self.args.command:
            cmd_buffer = b''
            while True:
                try:
                    while '\n' not in cmd_buffer.decode():
                        cmd_buffer += client_socket.recv(64)
                    result = self.execute(cmd_buffer.decode())
                    if result:
                        client_socket.send(result.encode())
                    cmd_buffer = b''
                except Exception as e:
                    print("Server killed {}".format(e))
                    client_socket.close()
                    sys.exit()

    def execute(self, cmd):
        cmd = cmd.strip()
        if not cmd:
            return
        result = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT)
        return result.decode()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Python Netcat Tool',
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     epilog=textwrap.dedent('''
                                     Example:
                                         netcat.py -t 192.168.1.1 -p 2222 -l -c #command shell
                                         netcat.py -t 192.168.1.1 -p 2222 -l -u=test.txt #upload file
                                         netcat.py -t 192.168.1.1 -p 2222 #connect to server
                                     '''))
    parser.add_argument('-c', '--command', action='store_true', help='open a command shell')
    parser.add_argument('-e', '--execute', help='execute specified command')
    parser.add_argument('-l', '--listen', action='store_true', help='open listening for server')
    parser.add_argument('-p', '--port', type=int, default=8080, help='input port')
    parser.add_argument('-t', '--target', default='127.0.0.1', help='input IP')
    parser.add_argument('-f', '--file', help='upload file')
    args = parser.parse_args()
    if args.listen:
        buffer = ''
    else:
        buffer = sys.stdin.read()

    nc = NetCat(args, buffer.encode())
    nc.run()