← 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-21 04:20:44 UTC
  • Revision ID: matt.giuca@gmail.com-20100721042044-uopyzuiuji6vlsu7
Project page: Moved instructions under 'latest submissions' above the table, or it will be hard to see them under hundreds of submissions.

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.
53
59
            os._exit(0) # kill off parent again.
54
60
        os.umask(077)
55
61
 
 
62
        try:
 
63
            MAXFD = os.sysconf("SC_OPEN_MAX")
 
64
        except:
 
65
            MAXFD = 256
 
66
 
 
67
        # Close all file descriptors, except the socket.
 
68
        for i in xrange(MAXFD):
 
69
            if i == s.fileno():
 
70
                continue
 
71
            try:
 
72
                os.close(i)
 
73
            except OSError:
 
74
                pass
 
75
 
 
76
        si = os.open(os.devnull, os.O_RDONLY)
 
77
        os.dup2(si, sys.stdin.fileno())
 
78
 
 
79
        so = os.open(os.devnull, os.O_WRONLY)
 
80
        os.dup2(so, sys.stdout.fileno())
 
81
 
 
82
        se = os.open(os.devnull, os.O_WRONLY)
 
83
        os.dup2(se, sys.stderr.fileno())
 
84
 
56
85
    if initializer:
57
86
        initializer()
58
87
 
59
88
    while True:
60
89
        (conn, addr) = s.accept()
 
90
        conn.settimeout(SOCKETTIMEOUT)
61
91
        try:
62
 
            # Grab the input
63
 
            buf = cStringIO.StringIO()
64
 
            blk = conn.recv(1024)
65
 
            while blk:
66
 
                buf.write(blk)
67
 
                try:
68
 
                    blk = conn.recv(1024, socket.MSG_DONTWAIT)
69
 
                except:
70
 
                    # Exception thrown if it WOULD block (but we
71
 
                    # told it not to wait) - ie. we are done
72
 
                    blk = None
73
 
            inp = buf.getvalue()
74
 
            env = cjson.decode(inp)
75
 
            
76
 
            # Check that the message is 
77
 
            digest = md5.new(env['content'] + magic).digest().encode('hex')
78
 
            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:
79
97
                conn.close()
80
98
                continue
81
99
 
82
 
            content = cjson.decode(env['content'])
83
 
 
84
100
            response = handler(content)
85
101
 
86
 
            conn.sendall(cjson.encode(response))
 
102
            send_netstring(conn, cjson.encode(response))
87
103
 
88
104
            conn.close()
89
105
 
90
106
        except Terminate, t:
91
107
            # Try and send final response and then terminate
92
108
            if t.final_response:
93
 
                conn.sendall(cjson.encode(t.final_response))
 
109
                send_netstring(conn, cjson.encode(t.final_response))
94
110
            conn.close()
95
111
            sys.exit(0)
96
112
        except Exception:
103
119
                "value": str(e_val),
104
120
                "traceback": tb_dump.getvalue()
105
121
            }
106
 
            conn.sendall(cjson.encode(json_exc))
 
122
            send_netstring(conn, cjson.encode(json_exc))
107
123
            conn.close()
108
124
 
109
125
 
110
126
def chat(host, port, msg, magic, decode = True):
111
127
    sok = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
112
128
    sok.connect((host, port))
113
 
    content = cjson.encode(msg)
114
 
    digest = md5.new(content + magic).digest().encode("hex")
115
 
    env = {'digest':digest,'content':content}
116
 
    sok.send(cjson.encode(env))
117
 
 
118
 
    buf = cStringIO.StringIO()
119
 
    blk = sok.recv(1024)
120
 
    while blk:
121
 
        buf.write(blk)
122
 
        try:
123
 
            blk = sok.recv(1024, socket.MSG_DONTWAIT)
124
 
        except socket.error, e:
125
 
            if e[0] != 11:
126
 
                raise
127
 
            # Exception thrown if it WOULD block (but we
128
 
            # told it not to wait) - ie. we are done
129
 
            blk = None
130
 
    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
 
131
136
    sok.close()
132
137
 
133
138
    if decode:
135
140
    else:
136
141
        return inp
137
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)