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)