← Back to branch summary

~azzar1/unity/add-show-desktop-key

« back to all changes in this revision

Viewing changes to ivle/chat.py

  • Committer: William Grant
  • Date: 2012-06-28 01:52:02 UTC
  • Revision ID: me@williamgrant.id.au-20120628015202-f6ru7o367gt6nvgz
Hah

Show diffs side-by-side

added added

removed removed

Lines of Context:
ivle/chat.py
19
19
# Author: Thomas Conway
20
20
# Date:   5/2/2008
21
21
 
22
 
import cjson
 
22
try:
 
23
    import json
 
24
except ImportError:
 
25
    import simplejson as json
 
26
 
23
27
import cStringIO
24
28
import hashlib
25
29
import sys
27
31
import socket
28
32
import traceback
29
33
 
 
34
SOCKETTIMEOUT = 60
 
35
BLOCKSIZE = 1024
 
36
 
30
37
class Terminate(Exception):
31
 
    """Exception thrown when server is to be shut down. It will attempt to sned 
32
 
    the final_response to the client and then exits"""
 
38
    """Exception thrown when server is to be shut down. It will attempt to
 
39
    send the final_response to the client and then exits"""
33
40
    def __init__(self, final_response=None):
34
41
        self.final_response = final_response
35
42
 
36
43
    def __str__(self):
37
44
        return repr(self.final_response)
38
45
 
 
46
class ProtocolError(Exception):
 
47
    """Exception thrown when client violates the the chat protocol"""
 
48
    pass
39
49
 
40
50
def start_server(port, magic, daemon_mode, handler, initializer = None):
41
51
    # Attempt to open the socket.
81
91
 
82
92
    while True:
83
93
        (conn, addr) = s.accept()
 
94
        conn.settimeout(SOCKETTIMEOUT)
84
95
        try:
85
 
            # Grab the input
86
 
            buf = cStringIO.StringIO()
87
 
            blk = conn.recv(1024)
88
 
            while blk:
89
 
                buf.write(blk)
90
 
                try:
91
 
                    blk = conn.recv(1024, socket.MSG_DONTWAIT)
92
 
                except:
93
 
                    # Exception thrown if it WOULD block (but we
94
 
                    # told it not to wait) - ie. we are done
95
 
                    blk = None
96
 
            inp = buf.getvalue()
97
 
            env = cjson.decode(inp)
98
 
            
99
 
            # Check that the message is 
100
 
            digest = hashlib.md5(env['content'] + magic).hexdigest()
101
 
            if env['digest'] != digest:
 
96
            # Grab the input and try to decode
 
97
            inp = recv_netstring(conn)
 
98
            try:
 
99
                content = decode(inp, magic)
 
100
            except ProtocolError:
102
101
                conn.close()
103
102
                continue
104
103
 
105
 
            content = cjson.decode(env['content'])
106
 
 
107
104
            response = handler(content)
108
105
 
109
 
            conn.sendall(cjson.encode(response))
 
106
            send_netstring(conn, json.dumps(response))
110
107
 
111
108
            conn.close()
112
109
 
113
110
        except Terminate, t:
114
111
            # Try and send final response and then terminate
115
112
            if t.final_response:
116
 
                conn.sendall(cjson.encode(t.final_response))
 
113
                send_netstring(conn, json.dumps(t.final_response))
117
114
            conn.close()
118
115
            sys.exit(0)
119
116
        except Exception:
126
123
                "value": str(e_val),
127
124
                "traceback": tb_dump.getvalue()
128
125
            }
129
 
            conn.sendall(cjson.encode(json_exc))
 
126
            send_netstring(conn, json.dumps(json_exc))
130
127
            conn.close()
131
128
 
132
129
 
133
130
def chat(host, port, msg, magic, decode = True):
134
131
    sok = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
135
132
    sok.connect((host, port))
136
 
    content = cjson.encode(msg)
137
 
    digest = hashlib.md5(content + magic).hexdigest()
138
 
    env = {'digest':digest,'content':content}
139
 
    sok.send(cjson.encode(env))
140
 
 
141
 
    buf = cStringIO.StringIO()
142
 
    blk = sok.recv(1024)
143
 
    while blk:
144
 
        buf.write(blk)
145
 
        try:
146
 
            blk = sok.recv(1024, socket.MSG_DONTWAIT)
147
 
        except socket.error, e:
148
 
            if e[0] != 11:
149
 
                raise
150
 
            # Exception thrown if it WOULD block (but we
151
 
            # told it not to wait) - ie. we are done
152
 
            blk = None
153
 
    inp = buf.getvalue()
 
133
    sok.settimeout(SOCKETTIMEOUT)
 
134
 
 
135
    out = encode(msg, magic)
 
136
 
 
137
    send_netstring(sok, out)
 
138
    inp = recv_netstring(sok)
 
139
 
154
140
    sok.close()
155
141
 
156
142
    if decode:
157
 
        return cjson.decode(inp)
 
143
        return json.loads(inp)
158
144
    else:
159
145
        return inp
160
146
 
 
147
def encode(message, magic):
 
148
    """Converts a message into a JSON serialisation and uses a magic
 
149
    string to attach a HMAC digest.
 
150
    """
 
151
    # XXX: Any reason that we double encode?
 
152
    content = json.dumps(message)
 
153
 
 
154
    digest = hashlib.md5(content + magic).hexdigest()
 
155
    env = {'digest':digest,'content':content}
 
156
    return json.dumps(env)
 
157
 
 
158
 
 
159
def decode(message, magic):
 
160
    """Takes a message with an attached HMAC digest and validates the message.
 
161
    """
 
162
    msg = json.loads(message)
 
163
 
 
164
    # Check that the message is valid
 
165
    digest = hashlib.md5(msg['content'] + magic).hexdigest()
 
166
    if msg['digest'] != digest:
 
167
        raise ProtocolError("HMAC digest is invalid")
 
168
    content = json.loads(msg['content'])
 
169
 
 
170
    return content
 
171
 
 
172
 
 
173
def send_netstring(sok, data):
 
174
    """ Sends a netstring to a socket
 
175
    """
 
176
    netstring = "%d:%s,"%(len(data),data)
 
177
    sok.sendall(netstring)
 
178
 
 
179
 
 
180
def recv_netstring(sok):
 
181
    """ Attempts to recieve a Netstring from a socket.
 
182
    Throws a ProtocolError if the received data violates the Netstring 
 
183
    protocol.
 
184
    """
 
185
    # Decode netstring
 
186
    size_buffer = []
 
187
    c = sok.recv(1)
 
188
    while c != ':':
 
189
        # Limit the Netstring to less than 10^10 bytes (~1GB).
 
190
        if len(size_buffer) >= 10:
 
191
            raise ProtocolError(
 
192
                    "Could not read Netstring size in first 9 bytes: '%s'"%(
 
193
                    ''.join(size_buffer)))
 
194
        size_buffer.append(c)
 
195
        c = sok.recv(1)
 
196
    try:
 
197
        # Message should be length plus ',' terminator
 
198
        recv_expected = int(''.join(size_buffer)) + 1
 
199
    except ValueError, e:
 
200
        raise ProtocolError("Could not decode Netstring size as int: '%s'"%(
 
201
                ''.join(size_buffer)))
 
202
 
 
203
    # Read data
 
204
    buf = []
 
205
    recv_data = sok.recv(min(recv_expected, BLOCKSIZE))
 
206
    recv_size = len(recv_data)
 
207
    while recv_size < recv_expected:
 
208
        buf.append(recv_data)
 
209
        recv_data = sok.recv(min(recv_expected-recv_size, 1024))
 
210
        recv_size = recv_size + len(recv_data)
 
211
    assert(recv_size == recv_expected)
 
212
 
 
213
    # Did we receive the correct amount?
 
214
    if recv_data[-1] != ',':
 
215
        raise ProtocolError("Netstring did not end with ','")
 
216
    buf.append(recv_data[:-1])
 
217
 
 
218
    return ''.join(buf)