fs/skeleton/usr/bin/upload-logs

#!/usr/bin/python
# Copyright 2011 Google Inc. All Rights Reserved.
#
"""A command-line tool for uploading to gfiber-dropbox.appspot.com."""

__author__ = 'apenwarr@google.com (Avery Pennarun)'


import os.path
import random
import signal
import StringIO
import sys
import time
import urllib
import zlib
import pycurl
import options

optspec = """
upload-logs [options...] <filenames...>
--
s,server=          The server URL [https://diag.cpe.gfsvc.com]
k,key=             Add a key/value pair (format "-k key=value")
stdin=             Also upload stdin, with the given virtual filename
"""

# Initial retry time for the exponential backoff retry loop.
# This will create a backoff retry, with times centered at:
# 30, 60, 120, 240, 480, 480
RETRY_INIT_DELAY = 30
RETRY_MAX_DELAY = 480


class HttpError(Exception):
  def __init__(self, status):
    self.status = status
    Exception.__init__(self, str(self))

  def __str__(self):
    return 'http status: %d' % (self.status)


def HttpDo(method, url, post_data=None, content_type=None):
  """Make an HTTPS request using pycurl and return the result."""
  proto, _ = urllib.splittype(url)
  assert proto.lower() in ('http', 'https')
  assert method in ('GET', 'POST')
  outdata = StringIO.StringIO()
  # The log upload server does not take kindly to Expect: 100-continue
  # so remove that.
  headers = ['User-Agent: upload-logs', 'Expect:']
  if content_type:
    headers.append('Content-Type: %s' % content_type)
  curl = pycurl.Curl()
  curl.setopt(pycurl.WRITEFUNCTION, outdata.write)
  curl.setopt(pycurl.FOLLOWLOCATION, 0)
  curl.setopt(pycurl.SSL_VERIFYPEER, 1)
  curl.setopt(pycurl.SSL_VERIFYHOST, 2)
  if os.path.exists('/etc/ssl/private/device.key'):
    curl.setopt(pycurl.SSLKEY, '/etc/ssl/private/device.key')
  if os.path.exists('/etc/ssl/certs/device.pem'):
    curl.setopt(pycurl.SSLCERT, '/etc/ssl/certs/device.pem')
  curl.setopt(pycurl.URL, url)
  curl.setopt(pycurl.HTTPHEADER, headers)
  if method == 'GET':
    curl.setopt(pycurl.HTTPGET, True)
  else:
    assert post_data is not None
    request_buffer = StringIO.StringIO(post_data)

    def Ioctl(cmd):
      if cmd == pycurl.IOCMD_RESTARTREAD:
        request_buffer.seek(0)
    curl.setopt(pycurl.POST, True)
    curl.setopt(pycurl.IOCTLFUNCTION, Ioctl)
    curl.setopt(pycurl.READFUNCTION, request_buffer.read)
    curl.setopt(pycurl.POSTFIELDSIZE, len(post_data))
  curl.perform()
  http_status = curl.getinfo(pycurl.HTTP_CODE)
  if http_status != 200:
    raise HttpError(http_status)
  curl.close()
  return outdata.getvalue()


def UploadFile(url, filename, fileobj, keys):
  """Actually upload the given file to the server."""
  while filename.startswith('/'):
    filename = filename[1:]
  start_url = os.path.join(url, 'upload', filename)
  if keys:
    start_url += '?' + urllib.urlencode(keys)
  upload_url = HttpDo('GET', start_url)

  splitter = 'foo-splitter-%f' % time.time()
  content_type = 'multipart/form-data; boundary=%s' % splitter

  basecontent = zlib.compress(fileobj.read())
  attachment = ('--%(splitter)s\r\n'
                'Content-Disposition: form-data; name="file";'
                '  filename="%(filename)s"\r\n'
                '\r\n'
                '%(data)s'
                '\r\n'
                '--%(splitter)s--\r\n'
                '\r\n'
                % dict(splitter=splitter,
                       filename=filename,
                       data=basecontent))

  # Retry upload forever until success.
  # Each iteration increase the delay which should give the server
  # more time to digest whatever data is has already received.
  i = 0
  while True:
    try:
      HttpDo('POST', upload_url, attachment, content_type)
    except HttpError, e:
      # This is the success case.
      if e.status == 302:
        break

      # If the server is overloaded, retry after some random delay.
      print "upload-logs failed: %s" % e.status

      # Retry interval is maximum of 5 minutes, with a random delay
      # of +/- 50% of the retry interval.
      delay = min(RETRY_MAX_DELAY, RETRY_INIT_DELAY * 2**i)
      rand_offset = random.uniform(-delay*0.5, delay*0.5)
      time.sleep(delay + rand_offset)
      i = min(i+1, 10)
    else:
      # http code 200 case.
      raise Exception('expected http response code 302')


def main():
  # set an alarm, in case our HTTP client (or anything else) hangs
  # for any reason
  signal.alarm(60)
  # Sending USR1 should now interrupt time.sleep()
  signal.signal(signal.SIGUSR1, lambda signum, frame: 0)

  o = options.Options(optspec)
  (opt, flags, extra) = o.parse(sys.argv[1:])  # pylint: disable-msg=W0612
  if not extra and not opt.stdin:
    o.fatal('at least one filename and/or --stdin expected')
  keys = []
  for k, v in flags:
    if k in ('-k', '--key'):
      keys.append(tuple(v.split('=', 1)))

  if opt.stdin:
    UploadFile(opt.server, opt.stdin, sys.stdin, keys)
  for filename in extra:
    UploadFile(opt.server, filename, open(filename), keys)


if __name__ == '__main__':
  main()