S3 transport for debian apt, using only python, boto and evil

s3.py

#!/usr/bin/env python -u
#
# The S3 Acquire Method
#
# Author:  Nathan Mehl <n@climate.com>
# Date:    2013-10-07
# License: GPLv2
#
# Copyright (c) 1999--2012 Red Hat, Inc.
# Copyright (c) 2013 The Climate Corporation
#
# Based on apt-transport-spacewalk by Simon Lukasik <xlukas08 [at] stud.fit.vutbr.cz>
#
# This software is licensed to you under the GNU General Public License,
# version 2 (GPLv2). There is NO WARRANTY for this software, express or
# implied, including the implied warranties of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. You should have received a copy of GPLv2
# along with this software; if not, see
# http://www.gnu.org/licenses/old-licenses/gpl-2.0.txt.


import hashlib
import re
import sys
import urllib
import warnings

warnings.filterwarnings("ignore", message="the md5 module is deprecated; use hashlib instead")

from urlparse import urlparse
from boto.s3.connection import S3Connection


class pkg_acquire_method:
    """
    This is slightly modified python variant of apt-pkg/acquire-method.
    It is a skeleton class that implements only very basic of apt methods
    functionality.
    """
    __eof = False

    def __init__(self):
        print "100 Capabilities\nVersion: 1.0\nSingle-Instance: true\n\n",

    def __get_next_msg(self):
        """
        Apt uses for communication with its methods the text protocol similar
        to http. This function parses the protocol messages from stdin.
        """
        if self.__eof:
            return None
        result = {};
        line = sys.stdin.readline()
        while line == '\n':
            line = sys.stdin.readline()
        if not line:
            self.__eof = True
            return None
        s = line.split(" ", 1)
        result['_number'] = int(s[0])
        result['_text'] = s[1].strip()

        while not self.__eof:
            line = sys.stdin.readline()
            if not line:
                self.__eof = True
                return result
            if line == '\n':
                return result
            s = line.split(":", 1)
            result[s[0]] = s[1].strip()

    def __dict2msg(self, msg):
        """Convert dictionary to http like message"""
        result = ""
        for item in msg.keys():
            if msg[item] != None:
                result += item + ": " + msg[item] + "\n"
        return result

    def status(self, **kwargs):
        print "102 Status\n%s\n" % self.__dict2msg(kwargs),

    def uri_start(self, msg):
        print "200 URI Start\n%s\n" % self.__dict2msg(msg),

    def uri_done(self, msg):
        print "201 URI Done\n%s\n" % self.__dict2msg(msg),

    def uri_failure(self, msg):
        print "400 URI Failure\n%s\n" % self.__dict2msg(msg),

    def run(self):
        """Loop through requests on stdin"""
        while True:
            msg = self.__get_next_msg()
            if msg == None:
                return 0
            if msg['_number'] == 600:
                try:
                    self.fetch(msg)
                except Exception, e:
                    self.fail(e.__class__.__name__ + ": " + str(e))
            else:
                return 100

    def fetch(self):
        raise NotImplementedError

    def fail(self):
        raise NotImplementedError


class s3_method(pkg_acquire_method):
    """
    S3 acquire method
    """
    bucket = None

    def fail(self, message = 's3 failure'):
        self.uri_failure({'URI': self.uri,
                          'Message': message})

    def _get_bucket(self):
        # deb s3://accesskeyid:[secretkey]@bucket/ lucid qa1 extras
        if self.bucket == None:
            s3_conn = S3Connection(self.uri_parsed.username,
                    urllib.unquote(self.uri_parsed.password))
            self.bucket = s3_conn.get_bucket(self.uri_parsed.hostname)

    def __urlparse(self, uri):
        """urllib.urlparse spits blood if there's a / in the password field"""
        lindex = uri.find('[')
        rindex = uri.find(']')
        secret_key = uri[lindex + 1:rindex]
        uri = uri.replace('[%s]' % secret_key, urllib.quote(secret_key, safe=''))
        return urlparse(uri)

    def fetch(self, msg):
        """
        Fetch the content from s3 bucket to the file.

        Acording to the apt protocol msg must contain: 'URI' and 'Filename'.
        Other possible keys are: 'Last-Modified', 'Index-File', 'Fail-Ignore'
        """
        self.uri = msg['URI']
        self.uri_parsed = self.__urlparse(msg['URI'])
        self.filename = msg['Filename']

        self._get_bucket()

        key = self.bucket.get_key(self.uri_parsed.path)
        self.status(URI = self.uri, Message = 'Checking key')

        if key == None:
            self.uri_failure({'URI': self.uri,
                              'Message': '404 key not found',
                              'FailReason': 'HttpError404'})
            return

        self.uri_start({'URI': self.uri,
                        'Size': str(key.size),
                        'Last-Modified': key.last_modified})

        f = open(self.filename, "w")
        hash_sha256 = hashlib.sha256()
        hash_md5 = hashlib.md5()
        while True:
            data = key.read(4096)
            if not len(data):
                break
            hash_sha256.update(data)
            hash_md5.update(data)
            f.write(data)
        key.close()
        f.close()

        self.uri_done({'URI': self.uri,
                       'Filename': self.filename,
                       'Size': str(key.size),
                       'Last-Modified': key.last_modified,
                       'MD5-Hash': hash_md5.hexdigest(),
                       'MD5Sum-Hash': hash_md5.hexdigest(),
                       'SHA256-Hash': hash_sha256.hexdigest()})


if __name__ == '__main__':
    try:
        method = s3_method()
        ret = method.run()
        sys.exit(ret)
    except KeyboardInterrupt:
        pass