← 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
 
import md5
 
28
import hashlib
25
29
import sys
26
30
import os
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 = md5.new(env['content'] + magic).digest().encode('hex')
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 = md5.new(content + magic).digest().encode("hex")
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)