pikesaku’s blog

個人的な勉強メモです。記載内容について一切の責任は持ちません。

Pythonのネットワークプラグラミング

netcatもどき(作成中)

# -*- coding: utf-8 -*-

import sys
import socket
import getopt
import threading
import subprocess
import argparse
import re
import time

CS          = 'utf-8'

# Client Mode Option
HOST        = None
PORT        = 0
SRC_PORT    = 0

# Server Mode Option
LISTEN      = False
LISTEN_PORT = 0
WRITE_FILE  = None
EXEC_FILE   = None
COMMAND     = None


def do_server():
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind(('0.0.0.0', LISTEN_PORT))
    server.listen(5)
    while True:
        client_socket, addr = server.accept()
        client_thread = threading.Thread(target=client_handler, args=(client_socket,))
        client_thread.start()


def client_handler(client_socket):
    if WRITE_FILE:
        file_buffer = b''

        while True:
            data = client_socket.recv(1024)
            # terminates session if OS X Client send Ctrl + D
            if data == b'\x04':
                break
            else:
                file_buffer += data

        try:
            # must open file with wb mode. Because data is Byte data.
            fd = open(WRITE_FILE, 'wb')
            fd.write(file_buffer)
            fd.close()
            m = 'Success: saved data to ' + WRITE_FILE + '\r\n'
        except:
            m = 'Failed: failed to save data to ' + WRITE_FILE + '\r\n'

        # Unicode(Str) -> Byte Data(UTF-8)
        client_socket.send(m.encode(CS))
        client_socket.close()
        



#    if EXECUTE:
#        output = run_command(EXECUTE)
#        client_socket.send(output)
#
#    if SHELL:
#        prompt = '(SHELL): '
#        client_socket.send(prompt)
#
#        while True:
#           cmd_buffer = None
#           while '\n' not in cmd_buffer:
#              cmd_budder += client_socket.recv(1024)


def do_client():
    print('making')
 

def get_opt():
    def valid_target(s):
        if not re.search('^\S+:[1-65536]$', s):
           raise argparse.ArgumentTypeError('Invalid value. Value must be HOST:PORT')
        return s

    parser = argparse.ArgumentParser(description='This is tool similar to netcat by python. -l or -t option is required')

    cl_mode = parser.add_argument_group('Run on Client Mode')
    cl_mode.add_argument('-t', dest='target',      help='host:port(REQUIRED)',   type=valid_target)
    cl_mode.add_argument('-p', dest='source_port', help='source port',           type=int)

    sv_mode = parser.add_argument_group('Run on Server Mode. -w, -e, -c options is exclusive')
    sv_mode.add_argument('-l', dest='listen_port', help='listen port(REQUIRED)', type=int)
    sv_mode.add_argument('-w', dest='write_file',  help='write data to file')
    cl_mode.add_argument('-e', dest='exec_file',   help='execute file')
    cl_mode.add_argument('-c', dest='command',     help='execute command')

    args = parser.parse_args()

    if args.target:
        #---------------------------------------------------------
        # Client Mode Option Set
        #---------------------------------------------------------
        if args.listen_port or args.write_file or args.exec_file or args.command:
            parser.print_help()
        global HOST
        global PORT
        global SRC_PORT
    
        HOST     = args.target.split(':')[0]
        PORT     = int(args.target.split(':')[1])
        SRC_PORT = args.source_port
    else:
        #---------------------------------------------------------
        # Server Mode Option Set
        #---------------------------------------------------------
        if args.source_port:
            parser.print_help()
        if not args.listen_port:
            parser.print_help()

        global LISTEN
        global LISTEN_PORT
        global WRITE_FILE
        global COMMAND
        global EXEC_FILE

        LISTEN = True
        LISTEN_PORT = args.listen_port
        WRITE_FILE  = args.write_file
        EXEC_FILE   = args.exec_file
        COMMAND     = args.command


def main():
    get_opt()
    if LISTEN:
        print('Server Mode Start')
        do_server() 
    else:
        print('Client Mode Start')
        do_client() 


if __name__ == '__main__':
    main()