利用Python实现简易版netcat

利用Python实现简易版netcat

from argparse import Action
import optparse
import sys
import socket
import threading
import subprocess
import shlex

class MyNetCat:
    def __init__(self) -> None:
        self.target = self.get_params()[0]
        self.port = self.get_params()[1]
        self.listen_mode = self.get_params()[2]
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            if self.listen_mode:                
                self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                self.socket.bind(('0.0.0.0',self.port))
                self.socket.listen(5)
            else:
                self.socket.connect((self.target, self.port))
        except Exception as e:
            print("Something is wrong: %s" % e)
            sys.exit()

    def get_params(self):
        parser = optparse.OptionParser('Usage:./%s -t target -p port -l listen mode ' % sys.argv[0])
        parser.add_option('-t', '--target', dest='target', type='string', help='Specify target IP address')
        parser.add_option('-p', '--port', dest='port', type='int', help='Specify port to connect or listen on')
        parser.add_option('-l', '--listen_mode', dest='listen_mode', action='store_true', help='Specify mode')
        options, args = parser.parse_args()
        if options.listen_mode:    #表明为监听模式，即服务器端
            listen_mode = True
        else:# 表明为客户端模式
            listen_mode = False
            if options.target is None:
                print("Specify target IP address to connect")
                sys.exit()
        if options.port is None:
            print("Specify port")
            sys.exit()
        return options.target, options.port, options.listen_mode
    
    def execute_command(self, command): #服务端接收客户端发送过来的命令，利用subprocess模块执行
        command = command.strip()
        try:
            if not command:
                return 
            output = subprocess.check_output(shlex.split(command), stderr=subprocess.STDOUT)
            return output.decode('utf-8')
        except:
            output = 'Failed to execute command'
            return output
    
    def client(self):  # 实现客户端功能
        while True:
            input_data = input('Enter command: ')
            if input_data.strip() == 'exit':
                break
            self.socket.send(input_data.encode('utf-8'))
            recv_data = self.socket.recv(4096)
            print(recv_data.decode('utf-8'))
    def client_handler(self, client_socket, client_addr):# 实现服务端针对每个连接的客户端进行处理
        while True:
            recv_data = client_socket.recv(4096)
            if len(recv_data) == 0:
                break
            output = self.execute_command(recv_data.decode('utf-8'))
            client_socket.send(output.encode('utf-8'))


    def server(self):# 实现服务器端功能
        while True:
            client_socket, client_addr = self.socket.accept()
            t = threading.Thread(target=self.client_handler, args=(client_socket, client_addr))
            t.start()
    
    def run(self):
        try:
            if self.listen_mode:
                self.server()
            else:
                self.client()
        except Exception as e:
            print("Failed to run: %s" % e)
            sys.exit()

if __name__ == '__main__':
    nc = MyNetCat()
    nc.run()