← Back to branch summary

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

« back to all changes in this revision

Viewing changes to ivle/chat.py

  • Committer: Matt Giuca
  • Date: 2010-07-27 06:53:03 UTC
  • Revision ID: matt.giuca@gmail.com-20100727065303-gs7fn3gc3ccaqfux
https://launchpad.net/bugs/610330
Changed database schema 'semester' table. 'year' and 'semester' fields now allow any length, not just 4 or 1 respectively. (LP: #610330).

Show diffs side-by-side

added added

removed removed

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