Add exmine mode
[naja.git] / pyweek_upload.py
1 '''
2 Upload script specifically engineered for the PyWeek challenge.
3
4 Handles authentication and gives upload progress feedback.
5 '''
6 import sys
7 import os
8 import httplib
9 import cStringIO
10 import socket
11 import time
12 import getopt
13
14
15 class Upload(object):
16     def __init__(self, filename):
17         self.filename = filename
18
19 boundary = '--------------GHSKFJDLGDS7543FJKLFHRE75642756743254'
20 sep_boundary = '\n--' + boundary
21 end_boundary = sep_boundary + '--'
22
23
24 def mimeEncode(data, sep_boundary=sep_boundary, end_boundary=end_boundary):
25     '''Take the mapping of data and construct the body of a
26     multipart/form-data message with it using the indicated boundaries.
27     '''
28     ret = cStringIO.StringIO()
29     for key, value in data.items():
30         # handle multiple entries for the same name
31         if not isinstance(value, list):
32             value = [value]
33         for value in value:
34             ret.write(sep_boundary)
35             if isinstance(value, Upload):
36                 ret.write('\nContent-Disposition: form-data; name="%s"' % key)
37                 filename = os.path.basename(value.filename)
38                 ret.write('; filename="%s"\n\n' % filename)
39                 value = open(os.path.join(value.filename), "rb").read()
40             else:
41                 ret.write('\nContent-Disposition: form-data; name="%s"' % key)
42                 ret.write("\n\n")
43                 value = str(value)
44             ret.write(str(value))
45             if value and value[-1] == '\r':
46                 ret.write('\n')  # write an extra newline
47     ret.write(end_boundary)
48     return ret.getvalue()
49
50
51 class Progress(object):
52     def __init__(self, info, data):
53         self.info = info
54         self.tosend = len(data)
55         self.total = self.tosend / 1024
56         self.data = cStringIO.StringIO(data)
57         self.start = self.now = time.time()
58         self.sent = 0
59         self.num = 0
60         self.stepsize = self.total / 100 or 1
61         self.steptimes = []
62         self.display()
63
64     def __iter__(self):
65         return self
66
67     def next(self):
68         self.num += 1
69         if self.sent >= self.tosend:
70             print self.info, 'done', ' ' * (75 - len(self.info) - 6)
71             sys.stdout.flush()
72             raise StopIteration
73
74         chunk = self.data.read(1024)
75         self.sent += len(chunk)
76         # print (self.num, self.stepsize, self.total, self.sent, self.tosend)
77
78         if self.num % self.stepsize:
79             return chunk
80         self.display()
81         return chunk
82
83     def display(self):
84         # figure how long we've spent - guess how long to go
85         now = time.time()
86         steptime = now - self.now
87         self.steptimes.insert(0, steptime)
88         if len(self.steptimes) > 5:
89             self.steptimes.pop()
90         steptime = sum(self.steptimes) / len(self.steptimes)
91         self.now = now
92         eta = steptime * ((self.total - self.num) / self.stepsize)
93
94         # tell it like it is (or might be)
95         if now - self.start > 3:
96             M = eta / 60
97             H = M / 60
98             M = M % 60
99             S = eta % 60
100             if self.total:
101                 s = '%s %2d%% (ETA %02d:%02d:%02d)' % (
102                     self.info, self.num * 100. / self.total, H, M, S)
103             else:
104                 s = '%s 0%% (ETA %02d:%02d:%02d)' % (self.info, H, M, S)
105         elif self.total:
106             s = '%s %2d%%' % (self.info, self.num * 100. / self.total)
107         else:
108             s = '%s %d done' % (self.info, self.num)
109         sys.stdout.write(s + ' ' * (75 - len(s)) + '\r')
110         sys.stdout.flush()
111
112
113 class progressHTTPConnection(httplib.HTTPConnection):
114     def progress_send(self, str):
115         """Send `str' to the server."""
116         if self.sock is None:
117             self.connect()
118
119         p = Progress('Uploading', str)
120         for chunk in p:
121             sent = 0
122             while sent != len(chunk):
123                 try:
124                     sent += self.sock.send(chunk)
125                 except socket.error, v:
126                     if v[0] == 32:      # Broken pipe
127                         self.close()
128                     raise
129                 p.display()
130
131
132 class progressHTTP(httplib.HTTP):
133     _connection_class = progressHTTPConnection
134
135     def _setup(self, conn):
136         httplib.HTTP._setup(self, conn)
137         self.progress_send = self._conn.progress_send
138
139
140 def http_request(data, server, port, url):
141     h = progressHTTP(server, port)
142
143     data = mimeEncode(data)
144     h.putrequest('POST', url)
145     h.putheader('Content-type', 'multipart/form-data; boundary=%s' % boundary)
146     h.putheader('Content-length', str(len(data)))
147     h.putheader('Host', server)
148     h.endheaders()
149
150     h.progress_send(data)
151
152     errcode, errmsg, headers = h.getreply()
153
154     f = h.getfile()
155     response = f.read().strip()
156     f.close()
157
158     print '%s %s' % (errcode, errmsg)
159     if response:
160         print response
161
162
163 def usage():
164     print '''This program is to be used to upload files to the PyWeek system.
165 You may use it to upload screenshots or code submissions.
166
167 REQUIRED ARGUMENTS:
168  -u   username
169  -p   password
170  -d   description of file
171  -c   file to upload
172  -e   entry short name
173
174 OPTIONAL ARGUMENTS:
175  -s   file is a screenshot
176  -f   file is FINAL submission
177  -h   override default host name (www.pyweek.org)
178  -P   override default host port (80)
179
180 In order to qualify for judging at the end of the challenge, you MUST
181 upload your source and check the "Final Submission" checkbox.
182 '''
183
184
185 if __name__ == '__main__':
186     try:
187         optlist, args = getopt.getopt(sys.argv[1:], 'e:u:p:sfd:h:P:c:')
188     except getopt.GetoptError, message:
189         print message
190         usage()
191         sys.exit(1)
192     host = 'www.pyweek.org'
193     port = 80
194     data = dict(version=2)
195     optional = {}
196     url = None
197     for opt, arg in optlist:
198         if opt == '-u':
199             data['user'] = arg
200         elif opt == '-p':
201             data['password'] = arg
202         elif opt == '-s':
203             optional['is_screenshot'] = 'yes'
204         elif opt == '-f':
205             optional['is_final'] = 'yes'
206         elif opt == '-d':
207             data['description'] = arg
208         elif opt == '-c':
209             data['content_file'] = Upload(arg)
210         elif opt == '-e':
211             url = '/e/%s/oup/' % arg
212         elif opt == '-h':
213             host = arg
214         elif opt == '-P':
215             port = int(arg)
216
217     if len(data) < 4 or url is None:
218         print 'Required argument missing'
219         usage()
220         sys.exit(1)
221
222     data.update(optional)
223     http_request(data, host, port, url)