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