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()
|