【脚本】使用Python对socat简单的控制

1
import argparse
2
import os
3
import re
4
import subprocess
5

6

7
def get_dict(d, key):
8
    if type(d) is not dict:
9
        raise Exception('not dict')
10
    dd = d.get(key)
11
    if dd is None:
12
        return ''
13
    return dd
14

15

16
def socat_port():
17
    output = subprocess.check_output("ps aux | grep socat", shell=True, universal_newlines=True)
18

19
    out = []
20
    for line in output.splitlines():
21
        match = re.compile(r'socat (\w+)-LISTEN:(\d+)').search(line)
22
        if not match:
23
            continue
24
        data = {
25
            'listen_protocol': match.group(1).upper(),
26
            'listen_port': match.group(2)
27
        }
28

29
        match = re.compile(r'(\w+):([^:]+):(\d+)$').search(line)
30
        if not match:
31
            continue
32
        data['remote_protocol'] = match.group(1).upper()
33
        data['remote_host'] = match.group(2)
34
        data['remote_port'] = match.group(3)
35

36
        match = re.compile(r'(\d+)').search(line)
37
        if not match:
38
            continue
39
        data['listen_pid'] = match.group(1)
40

41
        match = re.compile(r',range=([^,]+),').search(line)
42
        if match:
43
            data['listen_range'] = match.group(1)
44
        match = re.compile(r',bind=([^,]+),').search(line)
45
        if match:
46
            data['listen_bind'] = match.group(1)
47
        out.append(data)
48

49
    return out
50

51

52
if __name__ == '__main__':
53
    parser = argparse.ArgumentParser()
54
    parser.add_argument('-r', '--range', help='监听网段范围，例如：192.168.1.0/24')
55
    parser.add_argument('-b', '--bind', help='绑定监听地址，例如：192.168.1.1')
56
    parser.add_argument('-p', '--protocol', default='all', help='[all | tcp | udp]')
57
    parser.add_argument('positional', nargs='*',
58
                        help='<add | del | list> <本地端口> <远程地址>，例如：add 8000 192.168.1.2:8080')
59

60
    args = parser.parse_args()
61

62
    try:
63
        operation = args.positional[0]
64
        if operation != 'list':
65
            if len(args.positional) != 3:
66
                raise Exception()
67

68
            port = int(args.positional[1])
69
            raddr = args.positional[2]
70

71
            if not 0 < port < 1 << 16:
72
                print('端口不正确')
73
                exit(1)
74
    except:
75
        parser.print_usage()
76
        exit(1)
77

78
    if operation == 'add':
79
        command = 'socat {protocol}-LISTEN:{port},{option}fork,reuseaddr {protocol}:{raddr} &'
80
        option = ''
81
        if args.range is not None:
82
            option += 'range=%s,' % args.range
83
        if args.bind is not None:
84
            option += 'bind=%s,' % args.bind
85

86
        commands = []
87
        if args.protocol == 'all':
88
            commands.append({'protocol': 'TCP',
89
                             'command': command.format(protocol='TCP', port=port, option=option, raddr=raddr)})
90
            commands.append({'protocol': 'UDP',
91
                             'command': command.format(protocol='UDP', port=port, option=option, raddr=raddr)})
92
        elif args.protocol == 'tcp':
93
            commands.append({'protocol': 'TCP',
94
                             'command': command.format(protocol='TCP', port=port, option=option, raddr=raddr)})
95
        elif args.protocol == 'udp':
96
            commands.append({'protocol': 'UDP',
97
                             'command': command.format(protocol='UDP', port=port, option=option, raddr=raddr)})
98
        list = socat_port()
99
        for c in commands:
100
            bk = False
101
            for data in list:
102
                if data.get('listen_port') == str(c.get('port')) and data.get('listen_protocol') == c.get(
103
                        'protocol').upper():
104
                    bk = True
105
                    break
106
            if bk:
107
                print('{protocol} {port} 已被占用'.format(protocol=c.get('protocol'), port=c.get('port')))
108
                continue
109
            os.system(c.get('command'))
110
    elif operation == 'del':
111
        list = socat_port()
112
        command = 'kill %s'
113
        commands = []
114
        for data in list:
115
            if args.protocol != 'all' and data.get('listen_protocol') != args.protocol:
116
                continue
117
            if str(port) != data.get('listen_port'):
118
                continue
119
            commands.append(command % data.get('listen_pid'))
120

121
        for c in commands:
122
            os.system(c)
123
    elif operation == 'list':
124
        layout = '{listen_protocol}\t{listen_port}\t{listen_bind}\t{listen_range}\t{remote}'
125
        # print(layout.format(listen_protocol=''))
126
        for data in socat_port():
127
            print(layout.format(
128
                listen_protocol=get_dict(data, 'listen_protocol'),
129
                listen_port=get_dict(data, 'listen_port'),
130
                listen_bind=get_dict(data, 'listen_bind'),
131
                listen_range=get_dict(data, 'listen_range'),
132
                remote='%s:%s:%s' % (
133
                    get_dict(data, 'remote_protocol'), get_dict(data, 'remote_host'), get_dict(data, 'remote_port')),
134
            ))
135
    else:
136
        print('参数不正确')
137
        parser.print_usage()
138
        exit(1)