-
Notifications
You must be signed in to change notification settings - Fork 159
/
rpc.py
90 lines (80 loc) · 2.88 KB
/
rpc.py
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
"""
This is a RPC server/client implementation using socket
* It authenticate the client with a secret shared by client/server
* It uses JSON to serialize the function call payload
"""
import socket
import random, string
import hmac
import json
import threading
secret = "SECRET"
class RPCHandler:
def __init__(self, secret):
self._secret = secret
self._register = {}
def register_func(self, func):
self._register[func.__name__] = func
def handle_call(self, sock):
keymsg = ''.join([random.choice(string.lowercase) for i in range(8)])
sock.sendall(keymsg)
hash = hmac.new(self._secret, keymsg)
digest = hash.digest()
response = sock.recv(512)
if response != digest:
sock.sendall("Authentication Failed!")
sock.close()
else:
sock.sendall("Authenticated!")
try:
while True:
req = sock.recv(512)
d = json.loads(req)
funcname = d["name"]
args = d["args"]
kwargs = d["kwargs"]
print "Client calling %s(%s, %s)"%(funcname, args, kwargs)
try:
ret = self._register[funcname](*args, **kwargs)
sock.sendall(json.dumps({"ret": ret}))
except Exception as e:
sock.sendall(json.dumps({"exception": str(e)}))
except EOFError:
print "Closing RPC Handler..."
handler = RPCHandler(secret)
class RPCServer:
def __init__(self, address):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.bind(address)
def serve_forever(self):
self._sock.listen(0)
while True:
client_sock,_ = self._sock.accept()
thread = threading.Thread(target=handler.handle_call, args=(client_sock, ))
thread.daemon = True
thread.start()
class RPCProxy(object):
def __init__(self, address, secret):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.connect(address)
msg = self._sock.recv(512)
h = hmac.new(secret, msg)
self._sock.sendall(h.digest())
print self._sock.recv(512)
def __getattr__(self, name):
def proxy_func(*args, **kwargs):
payload = {
"name": name,
"args": args,
"kwargs": kwargs
}
self._sock.sendall(json.dumps(payload))
result = json.loads(self._sock.recv(512))
if "exception" in result:
raise Exception(result['exception'])
else:
return result['ret']
if name.startswith("_"):
return super(RPCProxy, self).__getattr__(name)
else:
return proxy_func