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