File size: 4,318 Bytes
345ee20
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import argparse
import socket
import itchat
import errno
import time
import threading

import os
import subprocess

parser = argparse.ArgumentParser(description="simple server!")
parser.add_argument('--mode', type=str)
parser.add_argument('--ip', type=str)
parser.add_argument('--port', type=int)
parser.add_argument('--timeout', type=int)

class MsgServer(object):
    def __init__(self, server_ip, server_port):
        self.init_chat()
        self.send('server chat logged in!')
        self.start_server(server_ip, server_port)

    def init_chat(self):
        itchat.auto_login(enableCmdQR=2)

    def send(self, msg, echo=True):
        if echo:
            print(msg)
        itchat.send(msg, toUserName='filehelper')

    def worker_thread(self, conn, addr):
        conn.settimeout(args.timeout)
        if conn is not None:
            self.send('job connected! [{}]'.format(addr))
        else:
            self.send('none connection!')
            return -1

        while(True):
            try:
                recv_data = conn.recv(1024)
            except socket.timeout as e:
                print('no msg...')
            else:
                msg_len = len(recv_data)
                if msg_len == 0:
                    self.send('connection break, waiting for other connections..')
                    break
                self.send(str(recv_data, encoding = 'utf-8'))
        conn.close()

    def start_server(self, server_ip, server_port):
        ip_port = (server_ip, server_port)
        s = socket.socket()
        s.bind(ip_port)

        # dump ip/port info to file
        with open('server.txt', 'w') as f:
            f.write('{} {}\n'.format(server_ip, server_port))

        s.listen()
        self.send('server listening on {}, waiting for job connection...'.format(server_ip))
        while(True):
            conn, addr = s.accept()
            threading.Thread(target=self.worker_thread, args=(conn, addr)).start()

class MsgClient(object):
    def __init__(self, server_ip, server_port):
        self._init_client(server_ip, server_port)
        self.send('I\'m client!\n')

    def send(self, msg, echo=True):
        self.s.send(bytes(msg, encoding = 'utf-8'))
        if echo:
            print(msg)

    def _init_client(self, server_ip, server_port):
        ip_port = (server_ip, server_port)
        self.s = socket.socket()
        self.s.connect(ip_port)

    def close(self):
        self.s.close()

def itchat_manager():
    def run_and_get(cmd, screen=False):
        process = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
        output = ''

        for line in process.stdout:
            line = line.decode('utf-8')
            if screen:
                print(line, end='', flush=True)
            output += line.strip(' ')

        return output

    @itchat.msg_register(itchat.content.TEXT)
    def text_reply(msg):
        res_txt = None
        cmd_dict = {
            'sq' : 'squeue -p VI_Face_V100',
            'sq1': 'squeue -p VI_Face_1080TI'
        }
        if msg.text in cmd_dict:
            cmd = cmd_dict[msg.text]
            res_txt = run_and_get(cmd)
        elif msg.text.startswith('exec:'):
            cmd = msg.text.replace('exec:', '')
            if os.system(cmd) == 0:
                res_txt = 'exec successed!'
            else:
                res_txt = 'exec failed!'
        elif msg.text.startswith('getinfo:'):
            cmd = msg.text.replace('getinfo:', '')
            res_txt = run_and_get(cmd)
        
        if res_txt is not None:
            itchat.send(res_txt, toUserName='filehelper')

    itchat.auto_login(enableCmdQR=2, hotReload=True)
    itchat.run(True)

if __name__ == '__main__':
    args = parser.parse_args()
    if args.mode == 'server':
        if args.ip is None or args.port is None:
            with open('server.txt', 'r') as f:
                line = f.read().strip().split()
                args.ip = line[0]
                args.port = int(line[1])
                print('reading ip & port from server.txt, {}:{}'.format(args.ip, args.port))
        s = MsgServer(args.ip, args.port)
    elif args.mode == 'manager':
        itchat_manager()
    else:
        s = MsgClient(args.ip, args.port)
        time.sleep(5)
        s.close()