<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Pyfiscan is free web-application vulnerability and version scanner and can be used to locate out-dated versions of common web-applications in Linux-servers. Example usa case is hosting-providers keeping eye on their users installations to keep up with security-updates. Supports content management systems, blogging softwares, image-galleries, version controlling programs, wikis, admin panels and bulletin boards. Fingerprints are easy to create and modify as user can write those in YAML-syntax.

@author Henri 'fgeek' Salo <henri@nerv.fi>
@copyright Copyright (c) 2009-2012 Henri Salo
@licence BSD

Known issues and/or bugs:
1: If instance is upgraded from Joomla 1.6.1 to 1.7.x by unzipping there will be both version files libraries/joomla/version.php and includes/version.php where first is the old one.

TODO: Should be a feature to run commands in detected installation directory e.g. if /home/example/public_html/ directory contains file php5.fcgi
TODO: Fingerprints to YAML and use decorators in functions. References: http://www.artima.com/weblogs/viewpost.jsp?thread=240808 http://www.python.org/dev/peps/pep-0318/ http://wiki.python.org/moin/PythonDecorators http://wiki.python.org/moin/PythonDecoratorLibrary
TODO: Argument --strip-output, which should remove homedir/startdir and location from output (stdin, csv and log)
TODO: If one fingerprint finds a match the process should finish and not be scanned with other fingerprints
TODO: There should be argument for looking specific programs in for example: -s joomla,smf
TODO: Add unittests
TODO: Add support to continue interrupted session (Tuomo Komulainen requested). Could be implemented using http://docs.python.org/library/atexit.html with knowledge of current working directory and queues
TODO: http://docs.python.org/library/functions.html#isinstance

Data in YAML-files could include following:
    CVE, CVSS2, OSVDB, Secunia
    Publication date, Fixed date
    CPE, ISS X-Force ID
    SecurityTracker Alert ID, Vendor URL
"""

try:
    import sys
    import time
    import logging
    import traceback
    import os
    from docopt import docopt
    from multiprocessing import Process, Queue, Pool
    from multiprocessing.util import log_to_stderr

    from database import Database
    from detect import yaml_fn_dict
    from file_helpers import \
        filepaths_in_dir, validate_directory, check_dir_execution_bit

    from issuereport import IssueReport, get_timestamp
except ImportError, error:
    print('Import error: %s' % error)
    sys.exit(1)

queue = Queue()
database = Database('yamls/')

def populate_directory(fargs):
    directory, checkmodes = fargs

    start_time = time.time()
    try:
        if not validate_directory(directory, checkmodes):
            return time.time() - start_time

        logging.debug('Populating: %s', directory)
        for filename in filepaths_in_dir(directory):
            for appname in database.issues:
                for loc in database.locations(appname, with_lists=False):
                    if filename.endswith(loc):
                        queue.put((filename, loc, appname))
    except Exception:
        logging.error(traceback.format_exc())

    return time.time() - start_time

def populate_userdir(fargs):
    predefined_locations = ['www', 'secure_www']
    userdir, checkmodes = fargs
    locations = []

    try:
        userdir = os.path.abspath(userdir)
        if not validate_directory(userdir, checkmodes):
            return locations

        public_html_location = userdir + '/public_html'
        if validate_directory(public_html_location, checkmodes):
            logging.debug('Appending to locations: %s', public_html_location)
            locations.append(public_html_location)

        sites_location = userdir + '/sites'
        if validate_directory(sites_location, checkmodes):
            for site in os.listdir(sites_location):
                sitedir = sites_location + '/' + site
                if checkmodes:
                    if not check_dir_execution_bit(sitedir):
                        continue

                for predefined_directory in predefined_locations:
                    sites_location_last = sitedir + '/' + predefined_directory
                    if validate_directory(sites_location_last, checkmodes):
                        logging.debug('Appending to locations: %s', sites_location_last)
                        locations.append(sites_location_last)
    except Exception:
        logging.error(traceback.format_exc())

    return locations

class PopulateScanQueue:
    def populate(self, directories, checkmodes=False):
        """ Populates worker queue for further scanning. Takes list of
            directories to be scanned and checkmodes boolean if execution bit should be
            taken into account. """
        try:
            """Use list of directories in loop to check if locations in data dictionary exists."""

            starttime = time.time()

            p = Pool()
            dirs = ((d, checkmodes) for d in directories)
            p.map(populate_directory, dirs, chunksize=200)

            # all done
            queue.put(None)

            logging.info('Scanning for locations finished. Elapsed time: %.4f', \
                         time.time() - starttime)

        except OSError:
            logging.error(traceback.format_exc())
            sys.exit(traceback.format_exc())
        except Exception:
            logging.error(traceback.format_exc())

    def populate_predefined(self, startdir, checkmodes):
        if not type(startdir) == str:
            sys.exit('Error in populate_predefined value startdir not a string. Value is: "%s" with type %s.' % (startdir, type(startdir)[0]))
        try:
            logging.debug('Populating predefined directories: %s', startdir)
            starttime = time.time()

            p = Pool()
            dirs = (startdir + '/' + d for d in os.listdir(startdir))
            udirs = p.imap_unordered(populate_userdir, \
                                     ((d, checkmodes) for d in dirs), \
                                     chunksize=200)
            p.close()
            locations = [item for sublist in udirs for item in sublist]

            logging.info('Total amount of locations: %s, time elapsed: %.4f', \
                         len(locations), time.time() - starttime)

            self.populate(locations, checkmodes)
        except Exception:
            logging.error(traceback.format_exc())

def compare_versions(secure_version, file_version, appname=None):
    """Comparison of found version numbers. Value current_version is predefined
       and file_version is found from file using grep. Value appname is used to
       separate different version numbering syntax"""
    try:
        if not type(secure_version) == str:
            logging.debug('Secure version must be a string when comparing: %s', \
                          secure_version)

        if not type(file_version) == str:
            logging.debug('Version from file must be a string when comparing: %s', \
                          file_version)

        if appname == 'WikkaWiki':  # Replace -p → .
            ver1 = secure_version.split('-')
            ver2 = file_version.split('-')
            secure_version = ver1[0] + '.' + ver1[1].lstrip('p')
            file_version = ver2[0] + '.' + ver2[1].lstrip('p')
        ver1 = secure_version.split('.')
        ver2 = file_version.split('.')
        ver1_bigger = 0
        for i in range(len(min(ver1, ver2))):
            if int(ver1[i]) == int(ver2[i]):
                pass
            elif int(ver1[i]) > int(ver2[i]):
                ver1_bigger = 1
                return ver1_bigger
            else:
                return ver1_bigger
    except Exception:
        logging.error(traceback.format_exc())

def handle_results(report, appname, file_version, item_location, application_cve, \
                   application_secure):
    try:
        logging.debug('%s with version %s from %s with vulnerability %s. This installation should be updated to at least version %s.', appname, file_version, item_location, application_cve, application_secure)
        print('%s Found: %s %s -> %s (%s)' % (get_timestamp(), item_location, file_version, application_secure, appname))
        report.add(appname, item_location, file_version, application_secure, application_cve)
    except Exception:
        logging.error(traceback.format_exc())

def Worker():
    """This is the actual worker which calls smaller functions in case of
    correct directory/file match is found.

        - Takes and removes item from queue
        - Detection in case of correct directory/file match is found
        - Compares found version against secure version in YAML
        - Calls logging

    Every worker runs in a loop.

    """
    try:
        report = IssueReport()
    except Exception:
        report.close()
        logging.error(traceback.format_exc())
        return

    while 1:
        try:
            item = queue.get()
            if not item:
                break

            item_location, location, appname = item
            logging.info('Processing: %s (%s)', appname, item_location)

            for issue in database.issues[appname].itervalues():
                logging.debug('Processing item %s with location %s with with appname %s issue %s', \
                              item_location, location, appname, issue)

                fn = yaml_fn_dict[issue['fingerprint']]
                file_version = fn(item_location, issue['regexp'])

                # Makes sure we don't go forward without version number from the file
                if file_version:
                    # Tests that version from file is smaller than secure version
                    # with application fingerprint-function
                    logging.debug('Comparing versions %s:%s for item %s', \
                                  issue['secure_version'], file_version, item_location)

                    if compare_versions(issue['secure_version'], file_version, appname):
                        # item_location is stripped from application location so that
                        # we get cleaner output and actual installation directory
                        install_dir = item_location[:item_location.find(location)]

                        # Calls result handler (goes to CSV and log)
                        handle_results(report, appname, file_version, install_dir, \
                                       issue['cve'], issue['secure_version'])
                else:
                    logging.debug('No version found from item: %s with regexp %s', \
                                  item_location, issue['regexp'])
        except Exception:
            logging.error(traceback.format_exc())
    report.close()

def main():
    logfile = 'pyfiscan.log'

    # Available logging levels, which are also hardcoded to usage
    levels = {'info': logging.INFO, 'debug': logging.DEBUG}

    usage = """
    Usage:
      pyfiscan.py
      pyfiscan.py -r <directory>
      pyfiscan.py --home <directory> [--check-modes]
      pyfiscan.py [-h|--help]
      pyfiscan.py --version

    Options:
      -r DIR            Scans directories recurssively.
      --home DIR        Specifies where the home-directories are located.
      --check-modes     Check using execution bit if we are allowed to traverse directories.
      -l LEVEL          Specifies logging level: info, debug

      If you do not spesify recursive-option predefined directories are scanned, which are:
        /home/user/sites/www
        /home/user/sites/secure-www
        /home/user/public_html

    """
    arguments = docopt(usage, version='pyfiscan 2.0')

    # Starttime is used to measure program runtime
    starttime = time.time()
    if arguments['-l']:
        level_name = arguments['-l']
    else:
        level_name = str('info')
    if not level_name in levels:
        print('No such log level. Available levels are: %s' % levels.keys())
        sys.exit(1)
    level = levels.get(level_name, logging.NOTSET)

    # Exit in case logfile is symlink
    if os.path.islink(logfile):
        sys.exit('Logfile %s is a symlink. Exiting..' % logfile)
    try:
        logging.basicConfig(filename=logfile, level=level, format='%(asctime)s %(levelname)s %(funcName)s:%(lineno)d %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
        os.chmod(logfile, 0600)
    except IOError as (errno, strerror):
        if errno == int('13'):
            sys.exit('Error while writing to logfile: %s' % strerror)
    try:
        # stderr to /dev/null
        devnull_fd = open(os.devnull, "w")
        sys.stderr = devnull_fd
        log_to_stderr()

        # Starts the asynchronous workers. Amount of workers is the same as cores in server.
        # http://docs.python.org/library/multiprocessing.html#multiprocessing.pool.multiprocessing.Pool
        logging.debug('Starting workers.')
        pool = Pool()
        pool.apply_async(Worker)

        # Starts the actual populator daemon to get possible locations, which will be verified by workers.
        # http://docs.python.org/library/multiprocessing.html#multiprocessing.Process
        logging.debug('Starting scan queue populator.')
        p = PopulateScanQueue()
        if arguments['-r']:
            logging.debug('Scanning recursively from path: %s', arguments['-r'])
            populator = Process(target=p.populate, args=([arguments['-r']],))
        elif arguments['-h']:
            logging.debug('Scanning predefined variables: %s', arguments['-h'])
            populator = Process(target=p.populate_predefined, args=(arguments['-h'], arguments['--check-modes'],))
        else:
            logging.debug('Scanning predefined variables: /home')
            populator = Process(target=p.populate_predefined, args=('/home', arguments['--check-modes'],))

        populator.start()
        """Prevents any more tasks from being submitted to the pool. Once all the tasks have been completed the worker processes exit using kill-signal None
        http://docs.python.org/library/multiprocessing.html#multiprocessing.pool.multiprocessing.Pool.close
        """
        populator.join()
        pool.close()
        pool.join()
        runtime = time.time() - starttime
        logging.info('Scanning ended, which took %s seconds', runtime)
    except KeyboardInterrupt:
        logging.info('Received keyboard interrupt. Exiting..')
        pool.join()
        populator.join()
        runtime = time.time() - starttime
        logging.info('Scanning ended, which took %s seconds', runtime)
    except Exception:
        logging.error(traceback.format_exc())

if __name__ == "__main__":
    main()
<commit_msg>I accidentaly. Debumb version :)<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Pyfiscan is free web-application vulnerability and version scanner and can be used to locate out-dated versions of common web-applications in Linux-servers. Example usa case is hosting-providers keeping eye on their users installations to keep up with security-updates. Supports content management systems, blogging softwares, image-galleries, version controlling programs, wikis, admin panels and bulletin boards. Fingerprints are easy to create and modify as user can write those in YAML-syntax.

@author Henri 'fgeek' Salo <henri@nerv.fi>
@copyright Copyright (c) 2009-2012 Henri Salo
@licence BSD

Known issues and/or bugs:
1: If instance is upgraded from Joomla 1.6.1 to 1.7.x by unzipping there will be both version files libraries/joomla/version.php and includes/version.php where first is the old one.

TODO: Should be a feature to run commands in detected installation directory e.g. if /home/example/public_html/ directory contains file php5.fcgi
TODO: Fingerprints to YAML and use decorators in functions. References: http://www.artima.com/weblogs/viewpost.jsp?thread=240808 http://www.python.org/dev/peps/pep-0318/ http://wiki.python.org/moin/PythonDecorators http://wiki.python.org/moin/PythonDecoratorLibrary
TODO: Argument --strip-output, which should remove homedir/startdir and location from output (stdin, csv and log)
TODO: If one fingerprint finds a match the process should finish and not be scanned with other fingerprints
TODO: There should be argument for looking specific programs in for example: -s joomla,smf
TODO: Add unittests
TODO: Add support to continue interrupted session (Tuomo Komulainen requested). Could be implemented using http://docs.python.org/library/atexit.html with knowledge of current working directory and queues
TODO: http://docs.python.org/library/functions.html#isinstance

Data in YAML-files could include following:
    CVE, CVSS2, OSVDB, Secunia
    Publication date, Fixed date
    CPE, ISS X-Force ID
    SecurityTracker Alert ID, Vendor URL
"""

try:
    import sys
    import time
    import logging
    import traceback
    import os
    from docopt import docopt
    from multiprocessing import Process, Queue, Pool
    from multiprocessing.util import log_to_stderr

    from database import Database
    from detect import yaml_fn_dict
    from file_helpers import \
        filepaths_in_dir, validate_directory, check_dir_execution_bit

    from issuereport import IssueReport, get_timestamp
except ImportError, error:
    print('Import error: %s' % error)
    sys.exit(1)

queue = Queue()
database = Database('yamls/')

def populate_directory(fargs):
    directory, checkmodes = fargs

    start_time = time.time()
    try:
        if not validate_directory(directory, checkmodes):
            return time.time() - start_time

        logging.debug('Populating: %s', directory)
        for filename in filepaths_in_dir(directory):
            for appname in database.issues:
                for loc in database.locations(appname, with_lists=False):
                    if filename.endswith(loc):
                        queue.put((filename, loc, appname))
    except Exception:
        logging.error(traceback.format_exc())

    return time.time() - start_time

def populate_userdir(fargs):
    predefined_locations = ['www', 'secure_www']
    userdir, checkmodes = fargs
    locations = []

    try:
        userdir = os.path.abspath(userdir)
        if not validate_directory(userdir, checkmodes):
            return locations

        public_html_location = userdir + '/public_html'
        if validate_directory(public_html_location, checkmodes):
            logging.debug('Appending to locations: %s', public_html_location)
            locations.append(public_html_location)

        sites_location = userdir + '/sites'
        if validate_directory(sites_location, checkmodes):
            for site in os.listdir(sites_location):
                sitedir = sites_location + '/' + site
                if checkmodes:
                    if not check_dir_execution_bit(sitedir):
                        continue

                for predefined_directory in predefined_locations:
                    sites_location_last = sitedir + '/' + predefined_directory
                    if validate_directory(sites_location_last, checkmodes):
                        logging.debug('Appending to locations: %s', sites_location_last)
                        locations.append(sites_location_last)
    except Exception:
        logging.error(traceback.format_exc())

    return locations

class PopulateScanQueue:
    def populate(self, directories, checkmodes=False):
        """ Populates worker queue for further scanning. Takes list of
            directories to be scanned and checkmodes boolean if execution bit should be
            taken into account. """
        try:
            """Use list of directories in loop to check if locations in data dictionary exists."""

            starttime = time.time()

            p = Pool()
            dirs = ((d, checkmodes) for d in directories)
            p.map(populate_directory, dirs, chunksize=200)

            # all done
            queue.put(None)

            logging.info('Scanning for locations finished. Elapsed time: %.4f', \
                         time.time() - starttime)

        except OSError:
            logging.error(traceback.format_exc())
            sys.exit(traceback.format_exc())
        except Exception:
            logging.error(traceback.format_exc())

    def populate_predefined(self, startdir, checkmodes):
        if not type(startdir) == str:
            sys.exit('Error in populate_predefined value startdir not a string. Value is: "%s" with type %s.' % (startdir, type(startdir)[0]))
        try:
            logging.debug('Populating predefined directories: %s', startdir)
            starttime = time.time()

            p = Pool()
            dirs = (startdir + '/' + d for d in os.listdir(startdir))
            udirs = p.imap_unordered(populate_userdir, \
                                     ((d, checkmodes) for d in dirs), \
                                     chunksize=200)
            p.close()
            locations = [item for sublist in udirs for item in sublist]

            logging.info('Total amount of locations: %s, time elapsed: %.4f', \
                         len(locations), time.time() - starttime)

            self.populate(locations, checkmodes)
        except Exception:
            logging.error(traceback.format_exc())

def compare_versions(secure_version, file_version, appname=None):
    """Comparison of found version numbers. Value current_version is predefined
       and file_version is found from file using grep. Value appname is used to
       separate different version numbering syntax"""
    try:
        if not type(secure_version) == str:
            logging.debug('Secure version must be a string when comparing: %s', \
                          secure_version)

        if not type(file_version) == str:
            logging.debug('Version from file must be a string when comparing: %s', \
                          file_version)

        if appname == 'WikkaWiki':  # Replace -p → .
            ver1 = secure_version.split('-')
            ver2 = file_version.split('-')
            secure_version = ver1[0] + '.' + ver1[1].lstrip('p')
            file_version = ver2[0] + '.' + ver2[1].lstrip('p')
        ver1 = secure_version.split('.')
        ver2 = file_version.split('.')
        ver1_bigger = 0
        for i in range(len(min(ver1, ver2))):
            if int(ver1[i]) == int(ver2[i]):
                pass
            elif int(ver1[i]) > int(ver2[i]):
                ver1_bigger = 1
                return ver1_bigger
            else:
                return ver1_bigger
    except Exception:
        logging.error(traceback.format_exc())

def handle_results(report, appname, file_version, item_location, application_cve, \
                   application_secure):
    try:
        logging.debug('%s with version %s from %s with vulnerability %s. This installation should be updated to at least version %s.', appname, file_version, item_location, application_cve, application_secure)
        print('%s Found: %s %s -> %s (%s)' % (get_timestamp(), item_location, file_version, application_secure, appname))
        report.add(appname, item_location, file_version, application_secure, application_cve)
    except Exception:
        logging.error(traceback.format_exc())

def Worker():
    """This is the actual worker which calls smaller functions in case of
    correct directory/file match is found.

        - Takes and removes item from queue
        - Detection in case of correct directory/file match is found
        - Compares found version against secure version in YAML
        - Calls logging

    Every worker runs in a loop.

    """
    try:
        report = IssueReport()
    except Exception:
        report.close()
        logging.error(traceback.format_exc())
        return

    while 1:
        try:
            item = queue.get()
            if not item:
                break

            item_location, location, appname = item
            logging.info('Processing: %s (%s)', appname, item_location)

            for issue in database.issues[appname].itervalues():
                logging.debug('Processing item %s with location %s with with appname %s issue %s', \
                              item_location, location, appname, issue)

                fn = yaml_fn_dict[issue['fingerprint']]
                file_version = fn(item_location, issue['regexp'])

                # Makes sure we don't go forward without version number from the file
                if file_version:
                    # Tests that version from file is smaller than secure version
                    # with application fingerprint-function
                    logging.debug('Comparing versions %s:%s for item %s', \
                                  issue['secure_version'], file_version, item_location)

                    if compare_versions(issue['secure_version'], file_version, appname):
                        # item_location is stripped from application location so that
                        # we get cleaner output and actual installation directory
                        install_dir = item_location[:item_location.find(location)]

                        # Calls result handler (goes to CSV and log)
                        handle_results(report, appname, file_version, install_dir, \
                                       issue['cve'], issue['secure_version'])
                else:
                    logging.debug('No version found from item: %s with regexp %s', \
                                  item_location, issue['regexp'])
        except Exception:
            logging.error(traceback.format_exc())
    report.close()

def main():
    logfile = 'pyfiscan.log'

    # Available logging levels, which are also hardcoded to usage
    levels = {'info': logging.INFO, 'debug': logging.DEBUG}

    usage = """
    Usage:
      pyfiscan.py
      pyfiscan.py -r <directory>
      pyfiscan.py --home <directory> [--check-modes]
      pyfiscan.py [-h|--help]
      pyfiscan.py --version

    Options:
      -r DIR            Scans directories recurssively.
      --home DIR        Specifies where the home-directories are located.
      --check-modes     Check using execution bit if we are allowed to traverse directories.
      -l LEVEL          Specifies logging level: info, debug

      If you do not spesify recursive-option predefined directories are scanned, which are:
        /home/user/sites/www
        /home/user/sites/secure-www
        /home/user/public_html

    """
    arguments = docopt(usage, version='pyfiscan 0.9')

    # Starttime is used to measure program runtime
    starttime = time.time()
    if arguments['-l']:
        level_name = arguments['-l']
    else:
        level_name = str('info')
    if not level_name in levels:
        print('No such log level. Available levels are: %s' % levels.keys())
        sys.exit(1)
    level = levels.get(level_name, logging.NOTSET)

    # Exit in case logfile is symlink
    if os.path.islink(logfile):
        sys.exit('Logfile %s is a symlink. Exiting..' % logfile)
    try:
        logging.basicConfig(filename=logfile, level=level, format='%(asctime)s %(levelname)s %(funcName)s:%(lineno)d %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
        os.chmod(logfile, 0600)
    except IOError as (errno, strerror):
        if errno == int('13'):
            sys.exit('Error while writing to logfile: %s' % strerror)
    try:
        # stderr to /dev/null
        devnull_fd = open(os.devnull, "w")
        sys.stderr = devnull_fd
        log_to_stderr()

        # Starts the asynchronous workers. Amount of workers is the same as cores in server.
        # http://docs.python.org/library/multiprocessing.html#multiprocessing.pool.multiprocessing.Pool
        logging.debug('Starting workers.')
        pool = Pool()
        pool.apply_async(Worker)

        # Starts the actual populator daemon to get possible locations, which will be verified by workers.
        # http://docs.python.org/library/multiprocessing.html#multiprocessing.Process
        logging.debug('Starting scan queue populator.')
        p = PopulateScanQueue()
        if arguments['-r']:
            logging.debug('Scanning recursively from path: %s', arguments['-r'])
            populator = Process(target=p.populate, args=([arguments['-r']],))
        elif arguments['-h']:
            logging.debug('Scanning predefined variables: %s', arguments['-h'])
            populator = Process(target=p.populate_predefined, args=(arguments['-h'], arguments['--check-modes'],))
        else:
            logging.debug('Scanning predefined variables: /home')
            populator = Process(target=p.populate_predefined, args=('/home', arguments['--check-modes'],))

        populator.start()
        """Prevents any more tasks from being submitted to the pool. Once all the tasks have been completed the worker processes exit using kill-signal None
        http://docs.python.org/library/multiprocessing.html#multiprocessing.pool.multiprocessing.Pool.close
        """
        populator.join()
        pool.close()
        pool.join()
        runtime = time.time() - starttime
        logging.info('Scanning ended, which took %s seconds', runtime)
    except KeyboardInterrupt:
        logging.info('Received keyboard interrupt. Exiting..')
        pool.join()
        populator.join()
        runtime = time.time() - starttime
        logging.info('Scanning ended, which took %s seconds', runtime)
    except Exception:
        logging.error(traceback.format_exc())

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>import abc
from functools import wraps
import logging
import struct

from amber.common import drivermsg_pb2

__author__ = 'paoolo'

LEN_SIZE = 2
LOGGER_NAME = 'Amber.Pipes'

class MessageHandler(object):
    def __init__(self, pipe_in, pipe_out):
        self.__amber_pipes = AmberPipes(self, pipe_in, pipe_out)

    def __call__(self, *args, **kwargs):
        self.__amber_pipes(*args, **kwargs)

    @abc.abstractmethod
    def handle_data_message(self, header, message):
        pass

    @abc.abstractmethod
    def handle_subscribe_message(self, header, message):
        pass

    @abc.abstractmethod
    def handle_unsubscribe_message(self, header, message):
        pass

    @abc.abstractmethod
    def handle_client_died_message(self, client_id):
        pass

    def get_pipes(self):
        return self.__amber_pipes

    @staticmethod
    def handle_and_response(func):
        @wraps(func)
        def wrapped(inst, received_header, received_message):
            response_header = drivermsg_pb2.DriverHdr()
            response_message = drivermsg_pb2.DriverMsg()

            response_message.type = drivermsg_pb2.DriverMsg.DATA
            response_message.ackNum = received_message.synNum

            response_header.clientIDs.extend(received_header.clientIDs)

            response_header, response_message = func(inst, received_header, received_message,
                                                     response_header, response_message)

            inst.get_pipes().write_header_and_message_to_pipe(response_header, response_message)

        return wrapped

class AmberPipes(object):
    def __init__(self, message_handler, pipe_in, pipe_out):
        self.__message_handler = message_handler
        self.__pipe_in, self.__pipe_out = pipe_in, pipe_out

        self.__logger = logging.Logger(LOGGER_NAME)
        self.__logger.addHandler(logging.StreamHandler())

    def __call__(self, *args, **kwargs):
        self.__logger.info('Pipes thread started.')
        self.__run_process()

    def __run_process(self):
        while True:
            header, message = self.__read_header_and_message_from_pipe()
            self.__handle_header_and_message(header, message)

    def __read_data_from_pipe(self, container):
        data = self.__read_and_unpack_data_from_pipe(LEN_SIZE)
        container.ParseFromString(data)
        return container

    def __read_header_and_message_from_pipe(self):
        """
        Read and parse header and message from pipe.

        :return: header and message
        """
        header = drivermsg_pb2.DriverHdr()
        message = drivermsg_pb2.DriverMsg()

        header = self.__read_data_from_pipe(header)
        message = self.__read_data_from_pipe(message)

        return header, message

    def __read_and_unpack_data_from_pipe(self, size):
        """
        Read and unpack data from pipe.

        :param size: size of length data
        :return: binary string
        """
        data = self.__read_from_pipe(size)
        # FIXME: can generate error, why?
        size = struct.unpack('!H', data)
        data = self.__read_from_pipe(size[0])
        return data

    def __read_from_pipe(self, size):
        """
        Read binary string from pipe.

        :param size: size of read string
        :return: binary string
        """
        return self.__pipe_in.read(size)

    def __handle_header_and_message(self, header, message):
        """
        Handle any message. Not serviced message are PONG and DRIVER_DIED.

        :param header: object of DriverHdr
        :param message: object of DriverMsg
        :return: nothing
        """
        if message.type == drivermsg_pb2.DriverMsg.DATA:
            self.__logger.debug('Received DATA message')
            self.__message_handler.handle_data_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.SUBSCRIBE:
            self.__logger.debug('Received SUBSCRIBE message')
            self.__message_handler.handle_subscribe_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.UNSUBSCRIBE:
            self.__logger.debug('Received UNSUBSCRIBE message')
            self.__message_handler.handle_unsubscribe_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.CLIENT_DIED:
            self.__logger.debug('Received CLIENT_DIED message')
            self.__handle_client_died_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.PING:
            self.__logger.debug('Received PING message')
            self.__handle_ping_message(header, message)

        else:
            self.__logger.warning('Received unknown type message, ignoring.')

    def __handle_client_died_message(self, header, _):
        """
        Handle CLIENT_DIED message which came from mediator.
        Handling message delegated to message handler.

        :param header: object of DriverHdr
        :return: nothing
        """
        if len(header.clientIDs) != 1:
            self.__logger.warning('CLIENT_DIED\'s clientID not set, ignoring.')

        else:
            self.__message_handler.handle_client_died_message(header.clientIDs[0])

    def __handle_ping_message(self, ping_header, ping_message):
        """
        Handle PING message which came from mediator.

        :param ping_header: object of DriverHdr
        :param ping_message: object of DriverMsg
        :return: nothing
        """
        if ping_message.HasField('synNum'):
            self.__logger.warning('PING\'s synNum not set, ignoring.')

        else:
            pong_message = drivermsg_pb2.DriverMsg()
            pong_message.type = drivermsg_pb2.DriverMsg.MsgType.PONG
            pong_message.ackNum = ping_message.synNum

            pong_header = drivermsg_pb2.DriverHdr()
            pong_header.clientIDs.extend(ping_header.clientIDs)

            self.__logger.debug('Send PONG message')

            self.write_header_and_message_to_pipe(pong_header, pong_message)

    def write_header_and_message_to_pipe(self, header, message):
        """
        Serialize and write header and message to pipe.

        :param header: object of DriverHdr
        :param message: object of DriverMsg
        :return: nothing
        """
        self.__logger.debug('Write header and message to pipe: header="%s", message="%s".' %
                            (str(header).strip(), str(message).strip()))

        header_data = header.SerializeToString()
        self.__pack_and_write_data_to_pipe(header_data)

        message_data = message.SerializeToString()
        self.__pack_and_write_data_to_pipe(message_data)

    def __pack_and_write_data_to_pipe(self, binary_data):
        """
        Pack and write data to pipe.

        :param binary_data: binary data as string
        :return: nothing
        """
        binary_data = struct.pack('!H', len(binary_data)) + binary_data
        self.__write_to_pipe(binary_data)

    def __write_to_pipe(self, binary_string):
        """
        Write string binary to pipe.

        :param binary_string: binary string
        :return: nothing
        """
        self.__pipe_out.write(binary_string)
        self.__pipe_out.flush()

<commit_msg>Use signed shorts in pack unpack<commit_after>import abc
from functools import wraps
import logging
import struct

from amber.common import drivermsg_pb2

__author__ = 'paoolo'

LEN_SIZE = 2
LOGGER_NAME = 'Amber.Pipes'

class MessageHandler(object):
    def __init__(self, pipe_in, pipe_out):
        self.__amber_pipes = AmberPipes(self, pipe_in, pipe_out)

    def __call__(self, *args, **kwargs):
        self.__amber_pipes(*args, **kwargs)

    @abc.abstractmethod
    def handle_data_message(self, header, message):
        pass

    @abc.abstractmethod
    def handle_subscribe_message(self, header, message):
        pass

    @abc.abstractmethod
    def handle_unsubscribe_message(self, header, message):
        pass

    @abc.abstractmethod
    def handle_client_died_message(self, client_id):
        pass

    def get_pipes(self):
        return self.__amber_pipes

    @staticmethod
    def handle_and_response(func):
        @wraps(func)
        def wrapped(inst, received_header, received_message):
            response_header = drivermsg_pb2.DriverHdr()
            response_message = drivermsg_pb2.DriverMsg()

            response_message.type = drivermsg_pb2.DriverMsg.DATA
            response_message.ackNum = received_message.synNum

            response_header.clientIDs.extend(received_header.clientIDs)

            response_header, response_message = func(inst, received_header, received_message,
                                                     response_header, response_message)

            inst.get_pipes().write_header_and_message_to_pipe(response_header, response_message)

        return wrapped

class AmberPipes(object):
    def __init__(self, message_handler, pipe_in, pipe_out):
        self.__message_handler = message_handler
        self.__pipe_in, self.__pipe_out = pipe_in, pipe_out

        self.__logger = logging.Logger(LOGGER_NAME)
        self.__logger.addHandler(logging.StreamHandler())

    def __call__(self, *args, **kwargs):
        self.__logger.info('Pipes thread started.')
        self.__run_process()

    def __run_process(self):
        while True:
            header, message = self.__read_header_and_message_from_pipe()
            self.__handle_header_and_message(header, message)

    def __read_data_from_pipe(self, container):
        data = self.__read_and_unpack_data_from_pipe(LEN_SIZE)
        container.ParseFromString(data)
        return container

    def __read_header_and_message_from_pipe(self):
        """
        Read and parse header and message from pipe.

        :return: header and message
        """
        header = drivermsg_pb2.DriverHdr()
        message = drivermsg_pb2.DriverMsg()

        header = self.__read_data_from_pipe(header)
        message = self.__read_data_from_pipe(message)

        return header, message

    def __read_and_unpack_data_from_pipe(self, size):
        """
        Read and unpack data from pipe.

        :param size: size of length data
        :return: binary string
        """
        data = self.__read_from_pipe(size)
        # FIXME: can generate error, why?
        size = struct.unpack('!h', data)
        data = self.__read_from_pipe(size[0])
        return data

    def __read_from_pipe(self, size):
        """
        Read binary string from pipe.

        :param size: size of read string
        :return: binary string
        """
        return self.__pipe_in.read(size)

    def __handle_header_and_message(self, header, message):
        """
        Handle any message. Not serviced message are PONG and DRIVER_DIED.

        :param header: object of DriverHdr
        :param message: object of DriverMsg
        :return: nothing
        """
        if message.type == drivermsg_pb2.DriverMsg.DATA:
            self.__logger.debug('Received DATA message')
            self.__message_handler.handle_data_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.SUBSCRIBE:
            self.__logger.debug('Received SUBSCRIBE message')
            self.__message_handler.handle_subscribe_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.UNSUBSCRIBE:
            self.__logger.debug('Received UNSUBSCRIBE message')
            self.__message_handler.handle_unsubscribe_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.CLIENT_DIED:
            self.__logger.debug('Received CLIENT_DIED message')
            self.__handle_client_died_message(header, message)

        elif message.type == drivermsg_pb2.DriverMsg.PING:
            self.__logger.debug('Received PING message')
            self.__handle_ping_message(header, message)

        else:
            self.__logger.warning('Received unknown type message, ignoring.')

    def __handle_client_died_message(self, header, _):
        """
        Handle CLIENT_DIED message which came from mediator.
        Handling message delegated to message handler.

        :param header: object of DriverHdr
        :return: nothing
        """
        if len(header.clientIDs) != 1:
            self.__logger.warning('CLIENT_DIED\'s clientID not set, ignoring.')

        else:
            self.__message_handler.handle_client_died_message(header.clientIDs[0])

    def __handle_ping_message(self, ping_header, ping_message):
        """
        Handle PING message which came from mediator.

        :param ping_header: object of DriverHdr
        :param ping_message: object of DriverMsg
        :return: nothing
        """
        if ping_message.HasField('synNum'):
            self.__logger.warning('PING\'s synNum not set, ignoring.')

        else:
            pong_message = drivermsg_pb2.DriverMsg()
            pong_message.type = drivermsg_pb2.DriverMsg.MsgType.PONG
            pong_message.ackNum = ping_message.synNum

            pong_header = drivermsg_pb2.DriverHdr()
            pong_header.clientIDs.extend(ping_header.clientIDs)

            self.__logger.debug('Send PONG message')

            self.write_header_and_message_to_pipe(pong_header, pong_message)

    def write_header_and_message_to_pipe(self, header, message):
        """
        Serialize and write header and message to pipe.

        :param header: object of DriverHdr
        :param message: object of DriverMsg
        :return: nothing
        """
        self.__logger.debug('Write header and message to pipe: header="%s", message="%s".' %
                            (str(header).strip(), str(message).strip()))

        header_data = header.SerializeToString()
        self.__pack_and_write_data_to_pipe(header_data)

        message_data = message.SerializeToString()
        self.__pack_and_write_data_to_pipe(message_data)

    def __pack_and_write_data_to_pipe(self, binary_data):
        """
        Pack and write data to pipe.

        :param binary_data: binary data as string
        :return: nothing
        """
        binary_data = struct.pack('!h', len(binary_data)) + binary_data
        self.__write_to_pipe(binary_data)

    def __write_to_pipe(self, binary_string):
        """
        Write string binary to pipe.

        :param binary_string: binary string
        :return: nothing
        """
        self.__pipe_out.write(binary_string)
        self.__pipe_out.flush()

<|endoftext|>


<commit_before>#!/usr/bin/env python

import argparse
import atexit
import logging
import sys
from datetime import datetime
from functools import partial, reduce
from pprint import pformat
from operator import itemgetter
from docker import Client
from humanfriendly import format_size

DEFAULT_DOCKER_BASE_URL = 'unix://var/run/docker.sock'
HELP_DOCKER_BASE_URL = (
    'Refers to the protocol+hostname+port where the '
    'Docker server is hosted. Defaults to %s') % DEFAULT_DOCKER_BASE_URL
DEFAULT_DOCKER_API_VERSION = 'auto'
HELP_DOCKER_API_VERSION = (
    'The version of the API the client will use. '
    'Defaults to use the API version provided by the server')
DEFAULT_DOCKER_HTTP_TIMEOUT = 5
HELP_DOCKER_HTTP_TIMEOUT = (
    'The HTTP request timeout, in seconds. '
    'Defaults to %d secs') % DEFAULT_DOCKER_HTTP_TIMEOUT
DEFAULT_IMAGES_TO_KEEP = 2
HELP_IMAGES_TO_KEEP = (
    'How many docker images to keep. '
    'Defaults to %d images') % DEFAULT_IMAGES_TO_KEEP
HELP_KEEP_NONE_IMAGES = 'Keep <none> images'
HELP_NOOP = 'Do nothing'
HELP_VERBOSE = 'Print images to delete'

def _exit():
    logging.shutdown()

def is_debug_on():
    return logging.getLogger().getEffectiveLevel() == logging.DEBUG

def debug_var(debug, name, var):
    if debug:
        logging.debug('Var %s has: %s' % (name, pformat(var)))

def setup_parser(parser):
    parser.add_argument('--debug', help='debug mode', action='store_true')
    parser.add_argument(
        '--base-url',
        help=HELP_DOCKER_BASE_URL,
        default=DEFAULT_DOCKER_BASE_URL)
    parser.add_argument(
        '--api-version',
        help=HELP_DOCKER_API_VERSION,
        default=DEFAULT_DOCKER_API_VERSION)
    parser.add_argument(
        '--http-timeout',
        help=HELP_DOCKER_HTTP_TIMEOUT,
        default=DEFAULT_DOCKER_HTTP_TIMEOUT,
        type=int)
    parser.add_argument(
        '--images-to-keep',
        help=HELP_IMAGES_TO_KEEP,
        default=DEFAULT_IMAGES_TO_KEEP,
        type=int)
    parser.add_argument(
        '--keep-none-images',
        help=HELP_KEEP_NONE_IMAGES,
        action='store_true')
    parser.add_argument('--noop', help=HELP_NOOP, action='store_true')
    parser.add_argument('--verbose', help=HELP_VERBOSE, action='store_true')
    return parser

def validate_args(args):
    if args.http_timeout < 0:
        sys.stderr.write('HTTP timeout should be 0 or bigger\n')
        sys.exit(1)
    if args.images_to_keep < 0:
        sys.stderr.write('Images to keep should be 0 or bigger\n')
        sys.exit(1)

def split_by_none(images_lists, dict_):
    if not isinstance(images_lists, tuple):
        raise TypeError('First argument should be a tuple')
    non_none_images, none_images = images_lists
    if u'<none>:<none>' in dict_[u'RepoTags']:
        none_images.append(dict_)
    else:
        non_none_images.append(dict_)
    return (non_none_images, none_images)

def split_images(images):
    return reduce(split_by_none, images, ([], []))

def remove_keys_from_dict(keys, dict_):
    return {k: v for k, v in dict_.iteritems() if k not in keys}

def add_image_to_grp_images(grp_images, image):
    repo, _ = image[u'RepoTags'][0].split(':')
    new_image = remove_keys_from_dict([u'RepoTags'], image)
    new_image[u'Tags'] = [e.split(':')[-1] for e in image[u'RepoTags']]
    if repo in grp_images:
        grp_images[repo].append(new_image)
    else:
        grp_images[repo] = [new_image]
    return grp_images

def group_by_repo(images):
    return reduce(add_image_to_grp_images, images, {})

def reverse_sort_images_created(images):
    return sorted(images, key=itemgetter(u'Created'), reverse=True)

def sort_images_in_repos(repos):
    return {k: reverse_sort_images_created(v) for k, v in repos.iteritems()}

def fix_none_image(image):
    new_image = remove_keys_from_dict([u'RepoTags'], image)
    new_image[u'Tags'] = image[u'RepoTags']
    return new_image

def beautify_image(image):
    new_image = remove_keys_from_dict(
        [u'RepoDigests', u'ParentId', u'Labels'],
        image)
    new_image[u'Created'] = datetime.fromtimestamp(
        image[u'Created']).isoformat(' ')
    new_image[u'Size'] = format_size(image[u'Size'])
    new_image[u'VirtualSize'] = format_size(image[u'VirtualSize'])
    return new_image

def print_images_to_delete(repos):
    print('Images to delete')
    print(pformat({k: [beautify_image(e) for e in v]
          for k, v in repos.iteritems()}))

def remove_docker_image(client, id_):
    try:
        client.remove_image(id_)
    except Exception:
        pass

def clean_images_in_repo(client, images):
    [remove_docker_image(client, image[u'Id']) for image in images]

def clean_repos(client, repos):
    [clean_images_in_repo(client, images) for images in repos.itervalues()]

def main():
    atexit.register(_exit)
    parser = setup_parser(argparse.ArgumentParser(
        description='Clean old docker images'))
    args = parser.parse_args()
    if args.debug:
        logging.basicConfig(level=logging.DEBUG)
    debug = partial(debug_var, debug=is_debug_on())
    debug(name='args', var=args)
    validate_args(args)
    client = Client(
        base_url=args.base_url,
        version=args.api_version,
        timeout=args.http_timeout)
    images = client.images()
    debug(name='images', var=images)
    non_none_images, none_images = split_images(images)
    debug(name='non_none_images', var=non_none_images)
    debug(name='none_images', var=none_images)
    repos = sort_images_in_repos(group_by_repo(non_none_images))
    debug(name='repos', var=repos)
    to_delete = {}
    if not args.keep_none_images:
        to_delete[u'<none>'] = [fix_none_image(e) for e in none_images]
    debug(name='to_delete', var=to_delete)
    repos_w_images = {k: v for k, v in repos.iteritems()
                      if len(v) > args.images_to_keep}
    debug(name='repos_w_images', var=repos_w_images)
    to_delete.update({k: v[args.images_to_keep:]
                      for k, v in repos_w_images.iteritems()})
    debug(name='to_delete', var=to_delete)
    if args.verbose:
        print_images_to_delete(to_delete)
    if args.noop:
        sys.exit(0)
    clean_repos(client, to_delete)

if __name__ == '__main__':
    main()
<commit_msg>Refactor validate_args<commit_after>#!/usr/bin/env python

import argparse
import atexit
import logging
import sys
from datetime import datetime
from functools import partial, reduce
from pprint import pformat
from operator import itemgetter
from docker import Client
from humanfriendly import format_size

DEFAULT_DOCKER_BASE_URL = 'unix://var/run/docker.sock'
HELP_DOCKER_BASE_URL = (
    'Refers to the protocol+hostname+port where the '
    'Docker server is hosted. Defaults to %s') % DEFAULT_DOCKER_BASE_URL
DEFAULT_DOCKER_API_VERSION = 'auto'
HELP_DOCKER_API_VERSION = (
    'The version of the API the client will use. '
    'Defaults to use the API version provided by the server')
DEFAULT_DOCKER_HTTP_TIMEOUT = 5
HELP_DOCKER_HTTP_TIMEOUT = (
    'The HTTP request timeout, in seconds. '
    'Defaults to %d secs') % DEFAULT_DOCKER_HTTP_TIMEOUT
DEFAULT_IMAGES_TO_KEEP = 2
HELP_IMAGES_TO_KEEP = (
    'How many docker images to keep. '
    'Defaults to %d images') % DEFAULT_IMAGES_TO_KEEP
HELP_KEEP_NONE_IMAGES = 'Keep <none> images'
HELP_NOOP = 'Do nothing'
HELP_VERBOSE = 'Print images to delete'

def _exit():
    logging.shutdown()

def is_debug_on():
    return logging.getLogger().getEffectiveLevel() == logging.DEBUG

def debug_var(debug, name, var):
    if debug:
        logging.debug('Var %s has: %s' % (name, pformat(var)))

def setup_parser(parser):
    parser.add_argument('--debug', help='debug mode', action='store_true')
    parser.add_argument(
        '--base-url',
        help=HELP_DOCKER_BASE_URL,
        default=DEFAULT_DOCKER_BASE_URL)
    parser.add_argument(
        '--api-version',
        help=HELP_DOCKER_API_VERSION,
        default=DEFAULT_DOCKER_API_VERSION)
    parser.add_argument(
        '--http-timeout',
        help=HELP_DOCKER_HTTP_TIMEOUT,
        default=DEFAULT_DOCKER_HTTP_TIMEOUT,
        type=int)
    parser.add_argument(
        '--images-to-keep',
        help=HELP_IMAGES_TO_KEEP,
        default=DEFAULT_IMAGES_TO_KEEP,
        type=int)
    parser.add_argument(
        '--keep-none-images',
        help=HELP_KEEP_NONE_IMAGES,
        action='store_true')
    parser.add_argument('--noop', help=HELP_NOOP, action='store_true')
    parser.add_argument('--verbose', help=HELP_VERBOSE, action='store_true')
    return parser

def validate_args(args):
    checks = [
        (lambda args: args.http_timeout < 0,
         'HTTP timeout should be 0 or bigger\n'),
        (lambda args: args.images_to_keep < 0,
         'Images to keep should be 0 or bigger\n')]
    if [sys.stderr.write(msg) for checker, msg in checks if checker(args)]:
        sys.exit(1)

def split_by_none(images_lists, dict_):
    if not isinstance(images_lists, tuple):
        raise TypeError('First argument should be a tuple')
    non_none_images, none_images = images_lists
    if u'<none>:<none>' in dict_[u'RepoTags']:
        none_images.append(dict_)
    else:
        non_none_images.append(dict_)
    return (non_none_images, none_images)

def split_images(images):
    return reduce(split_by_none, images, ([], []))

def remove_keys_from_dict(keys, dict_):
    return {k: v for k, v in dict_.iteritems() if k not in keys}

def add_image_to_grp_images(grp_images, image):
    repo, _ = image[u'RepoTags'][0].split(':')
    new_image = remove_keys_from_dict([u'RepoTags'], image)
    new_image[u'Tags'] = [e.split(':')[-1] for e in image[u'RepoTags']]
    if repo in grp_images:
        grp_images[repo].append(new_image)
    else:
        grp_images[repo] = [new_image]
    return grp_images

def group_by_repo(images):
    return reduce(add_image_to_grp_images, images, {})

def reverse_sort_images_created(images):
    return sorted(images, key=itemgetter(u'Created'), reverse=True)

def sort_images_in_repos(repos):
    return {k: reverse_sort_images_created(v) for k, v in repos.iteritems()}

def fix_none_image(image):
    new_image = remove_keys_from_dict([u'RepoTags'], image)
    new_image[u'Tags'] = image[u'RepoTags']
    return new_image

def beautify_image(image):
    new_image = remove_keys_from_dict(
        [u'RepoDigests', u'ParentId', u'Labels'],
        image)
    new_image[u'Created'] = datetime.fromtimestamp(
        image[u'Created']).isoformat(' ')
    new_image[u'Size'] = format_size(image[u'Size'])
    new_image[u'VirtualSize'] = format_size(image[u'VirtualSize'])
    return new_image

def print_images_to_delete(repos):
    print('Images to delete')
    print(pformat({k: [beautify_image(e) for e in v]
          for k, v in repos.iteritems()}))

def remove_docker_image(client, id_):
    try:
        client.remove_image(id_)
    except Exception:
        pass

def clean_images_in_repo(client, images):
    [remove_docker_image(client, image[u'Id']) for image in images]

def clean_repos(client, repos):
    [clean_images_in_repo(client, images) for images in repos.itervalues()]

def main():
    atexit.register(_exit)
    parser = setup_parser(argparse.ArgumentParser(
        description='Clean old docker images'))
    args = parser.parse_args()
    if args.debug:
        logging.basicConfig(level=logging.DEBUG)
    debug = partial(debug_var, debug=is_debug_on())
    debug(name='args', var=args)
    validate_args(args)
    client = Client(
        base_url=args.base_url,
        version=args.api_version,
        timeout=args.http_timeout)
    images = client.images()
    debug(name='images', var=images)
    non_none_images, none_images = split_images(images)
    debug(name='non_none_images', var=non_none_images)
    debug(name='none_images', var=none_images)
    repos = sort_images_in_repos(group_by_repo(non_none_images))
    debug(name='repos', var=repos)
    to_delete = {}
    if not args.keep_none_images:
        to_delete[u'<none>'] = [fix_none_image(e) for e in none_images]
    debug(name='to_delete', var=to_delete)
    repos_w_images = {k: v for k, v in repos.iteritems()
                      if len(v) > args.images_to_keep}
    debug(name='repos_w_images', var=repos_w_images)
    to_delete.update({k: v[args.images_to_keep:]
                      for k, v in repos_w_images.iteritems()})
    debug(name='to_delete', var=to_delete)
    if args.verbose:
        print_images_to_delete(to_delete)
    if args.noop:
        sys.exit(0)
    clean_repos(client, to_delete)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>import pysb.core
from pysb.generator.bng import BngGenerator
import os
import subprocess
import random
import re
import itertools
import sympy
import numpy
from StringIO import StringIO

# Cached value of BNG path
_bng_path = None

def _get_bng_path():
    """Return the path to BioNetGen's BNG2.pl, based on either the BNGHOME
    environment variable if it's set, or a few hard-coded standard locations."""

    global _bng_path

    # Just return cached value if it's available
    if _bng_path:
        return _bng_path

    path_var = 'BNGHOME'
    dist_dirs = [
        '/usr/local/share/BioNetGen',
        'c:/Program Files/BioNetGen',
        ]
    # BNG 2.1.8 moved BNG2.pl up out of the Perl2 subdirectory, so to be more
    # compatible we check both the old and new locations.
    script_subdirs = ['', 'Perl2']

    def check_dist_dir(dist_dir):
        # Return the full path to BNG2.pl inside a BioNetGen distribution
        # directory, or False if directory does not contain a BNG2.pl in one of
        # the expected places.
        for subdir in script_subdirs:
            script_path = os.path.join(dist_dir, subdir, 'BNG2.pl')
            if os.access(script_path, os.F_OK):
                return script_path
        else:
            return False

    # First check the environment variable, which has the highest precedence
    if path_var in os.environ:
        script_path = check_dist_dir(os.environ[path_var])
        if not script_path:
            raise Exception('Environment variable %s is set but BNG2.pl could'
                            ' not be found there' % path_var)
    # If the environment variable isn't set, check the standard locations
    else:
        for dist_dir in dist_dirs:
            script_path = check_dist_dir(dist_dir)
            if script_path:
                break
        else:
            raise Exception('Could not find BioNetGen installed in one of the '
                            'following locations:' +
                            ''.join('\n    ' + d for d in dist_dirs))
    # Cache path for future use
    _bng_path = script_path
    return script_path

class GenerateNetworkError(RuntimeError):
    pass

generate_network_code = """
begin actions
generate_network({overwrite=>1});
end actions
"""

def _parse_bng_outfile(out_filename):
    """Parse a .gdat or .cdat outputfile produced by the ODE or SSA
       algorithm run by BNG.
       The format of the output files is an initial line of the form
       #   time   obs1    obs2    obs3  ...
       The column headers are separated by a differing number of spaces
       (not tabs). This function parses out the column names and then
       uses the numpy.loadtxt method to load the outputfile into a
       numpy record array.
    """

    try:
        out_file = open(out_filename, 'r')

        line = out_file.readline().strip() # Get the first line
        out_file.close()
        line = line[2:]  # strip off opening '# '
        raw_names = re.split('\s*', line)
        column_names = [raw_name for raw_name in raw_names if not raw_name == '']

        # Create the dtype argument for the numpy record array
        dt = zip(column_names, ('float',)*len(column_names))

        # Load the output file as a numpy record array, skip the name row
        arr = numpy.loadtxt(out_filename, dtype=dt, skiprows=1)

    except Exception as e:
        # FIXME special Exception/Error?
        raise Exception("problem parsing BNG outfile: " + str(e))

    return arr

def run_ssa(model, t_end=10, n_steps=100, output_dir='/tmp', cleanup=True):
    """Run a model through BNG's SSA simulator and return
    the simulation data as a numpy record array.
    """

    run_ssa_code = """
    begin actions
    generate_network({overwrite=>1});
    simulate_ssa({t_end=>%f, n_steps=>%d});\n
    end actions
    """ % (t_end, n_steps)

    gen = BngGenerator(model)
    bng_filename = '%s_%d_%d_temp.bngl' % (model.name,
                            os.getpid(), random.randint(0, 10000))
    gdat_filename = bng_filename.replace('.bngl', '.gdat')
    cdat_filename = bng_filename.replace('.bngl', '.cdat')
    net_filename = bng_filename.replace('.bngl', '.net')

    output = StringIO()

    try:
        #import pdb; pdb.set_trace()
        working_dir = os.getcwd()
        os.chdir(output_dir)
        bng_file = open(bng_filename, 'w')
        bng_file.write(gen.get_content())
        bng_file.write(run_ssa_code)
        bng_file.close()
        p = subprocess.Popen(['perl', _get_bng_path(), bng_filename],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (p_out, p_err) = p.communicate()
        if p.returncode:
            raise GenerateNetworkError(p_out.rstrip("at line")+"\n"+p_err.rstrip())

        output_arr = _parse_bng_outfile(gdat_filename)
        #ssa_file = open(ssa_filename, 'r')
        #output.write(ssa_file.read())
        #net_file.close()
        #if append_stdout:
        #    output.write("#\n# BioNetGen execution log follows\n# ==========")
        #    output.write(re.sub(r'(^|\n)', r'\n# ', p_out))
    finally:
        if cleanup:
            for filename in [bng_filename, gdat_filename,
                             cdat_filename, net_filename]:
                if os.access(filename, os.F_OK):
                    os.unlink(filename)
        os.chdir(working_dir)
    return output_arr

def generate_network(model, cleanup=True, append_stdout=False):
    """Run a model through BNG's generate_network function and return
    the content of the resulting .net file as a string"""
    gen = BngGenerator(model)
    bng_filename = '%s_%d_%d_temp.bngl' % (model.name, os.getpid(), random.randint(0, 10000))
    net_filename = bng_filename.replace('.bngl', '.net')
    output = StringIO()
    try:
        bng_file = open(bng_filename, 'w')
        bng_file.write(gen.get_content())
        bng_file.write(generate_network_code)
        bng_file.close()
        p = subprocess.Popen(['perl', _get_bng_path(), bng_filename],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (p_out, p_err) = p.communicate()
        if p.returncode:
            raise GenerateNetworkError(p_out.rstrip()+"\n"+p_err.rstrip())
        net_file = open(net_filename, 'r')
        output.write(net_file.read())
        net_file.close()
        if append_stdout:
            output.write("#\n# BioNetGen execution log follows\n# ==========")
            output.write(re.sub(r'(^|\n)', r'\n# ', p_out))
    finally:
        if cleanup:
            for filename in [bng_filename, net_filename]:
                if os.access(filename, os.F_OK):
                    os.unlink(filename)
    return output.getvalue()

def generate_equations(model):
    # only need to do this once
    # TODO track "dirty" state, i.e. "has model been modified?"
    #   or, use a separate "math model" object to contain ODEs
    if model.odes:
        return
    lines = iter(generate_network(model).split('\n'))
    try:
        while 'begin species' not in lines.next():
            pass
        model.species = []
        while True:
            line = lines.next()
            if 'end species' in line: break
            _parse_species(model, line)

        while 'begin reactions' not in lines.next():
            pass
        model.odes = [sympy.S(0)] * len(model.species)
        reaction_cache = {}
        while True:
            line = lines.next()
            if 'end reactions' in line: break
            (number, reactants, products, rate, rule) = line.strip().split()
            # the -1 is to switch from one-based to zero-based indexing
            reactants = tuple(int(r) - 1 for r in reactants.split(','))
            products = tuple(int(p) - 1 for p in products.split(','))
            rate = rate.rsplit('*')
            (rule_name, is_reverse) = re.match(r'#(\w+)(?:\((reverse)\))?', rule).groups()
            is_reverse = bool(is_reverse)
            r_names = ['s%d' % r for r in reactants]
            combined_rate = sympy.Mul(*[sympy.S(t) for t in r_names + rate])
            rule = model.rules[rule_name]
            reaction = {
                'reactants': reactants,
                'products': products,
                'rate': combined_rate,
                'rule': rule_name,
                'reverse': is_reverse,
                }
            model.reactions.append(reaction)
            key = (rule_name, reactants, products)
            key_reverse = (rule_name, products, reactants)
            reaction_bd = reaction_cache.get(key_reverse)
            if reaction_bd is None:
                # make a copy of the reaction dict
                reaction_bd = dict(reaction)
                # default to false until we find a matching reverse reaction
                reaction_bd['reversible'] = False
                reaction_cache[key] = reaction_bd
                model.reactions_bidirectional.append(reaction_bd)
            else:
                reaction_bd['reversible'] = True
                reaction_bd['rate'] -= combined_rate
            for p in products:
                model.odes[p] += combined_rate
            for r in reactants:
                model.odes[r] -= combined_rate
        # fix up reactions whose reverse version we saw first
        for r in model.reactions_bidirectional:
            if r['reverse']:
                r['reactants'], r['products'] = r['products'], r['reactants']
                r['rate'] *= -1
            # now the 'reverse' value is no longer needed
            del r['reverse']

        while 'begin groups' not in lines.next():
            pass
        while True:
            line = lines.next()
            if 'end groups' in line: break
            _parse_group(model, line)

    except StopIteration as e:
        pass

def _parse_species(model, line):
    index, species, value = line.strip().split()
    complex_compartment_name, complex_string = re.match(r'(?:@(\w+)::)?(.*)', species).groups()
    monomer_strings = complex_string.split('.')
    monomer_patterns = []
    for ms in monomer_strings:
        monomer_name, site_strings, monomer_compartment_name = re.match(r'(\w+)\(([^)]*)\)(?:@(\w+))?', ms).groups()
        site_conditions = {}
        if len(site_strings):
            for ss in site_strings.split(','):
                # FIXME this should probably be done with regular expressions
                if '!' in ss and '~' in ss:
                    site_name, condition = ss.split('~')
                    state, bond = condition.split('!')
                    if bond == '?':
                        bond = pysb.core.WILD
                    elif bond == '!':
                        bond = pysb.core.ANY
                    else:
                        bond = int(bond)
                    condition = (state, bond)
                elif '!' in ss:
                    site_name, condition = ss.split('!', 1)
                    if '!' in condition:
                        condition = [int(c) for c in condition.split('!')]
                    else:
                        condition = int(condition)
                elif '~' in ss:
                    site_name, condition = ss.split('~')
                else:
                    site_name, condition = ss, None
                site_conditions[site_name] = condition
        monomer = model.monomers[monomer_name]
        monomer_compartment = model.compartments.get(monomer_compartment_name)
        mp = pysb.core.MonomerPattern(monomer, site_conditions, monomer_compartment)
        monomer_patterns.append(mp)

    complex_compartment = model.compartments.get(complex_compartment_name)
    cp = pysb.core.ComplexPattern(monomer_patterns, complex_compartment)
    model.species.append(cp)

def _parse_group(model, line):
    # values are number (which we ignore), name, and species list
    values = line.strip().split()
    obs = model.observables[values[1]]
    if len(values) == 3:
        # combination is a comma separated list of [coeff*]speciesnumber
        for product in values[2].split(','):
            terms = product.split('*')
            # if no coeff given (just species), insert a coeff of 1
            if len(terms) == 1:
                terms.insert(0, 1)
            obs.coefficients.append(int(terms[0]))
            # -1 to change to 0-based indexing
            obs.species.append(int(terms[1]) - 1)
<commit_msg>improved documentation on bng module<commit_after>import pysb.core
from pysb.generator.bng import BngGenerator
import os
import subprocess
import random
import re
import itertools
import sympy
import numpy
from StringIO import StringIO

# Cached value of BNG path
_bng_path = None

def _get_bng_path():
    """
    Return the path to BioNetGen's BNG2.pl.

    Looks for a BNG distribution at the path stored in the BNGHOME environment
    variable if that's set, or else in a few hard-coded standard locations.

    """

    global _bng_path

    # Just return cached value if it's available
    if _bng_path:
        return _bng_path

    path_var = 'BNGHOME'
    dist_dirs = [
        '/usr/local/share/BioNetGen',
        'c:/Program Files/BioNetGen',
        ]
    # BNG 2.1.8 moved BNG2.pl up out of the Perl2 subdirectory, so to be more
    # compatible we check both the old and new locations.
    script_subdirs = ['', 'Perl2']

    def check_dist_dir(dist_dir):
        # Return the full path to BNG2.pl inside a BioNetGen distribution
        # directory, or False if directory does not contain a BNG2.pl in one of
        # the expected places.
        for subdir in script_subdirs:
            script_path = os.path.join(dist_dir, subdir, 'BNG2.pl')
            if os.access(script_path, os.F_OK):
                return script_path
        else:
            return False

    # First check the environment variable, which has the highest precedence
    if path_var in os.environ:
        script_path = check_dist_dir(os.environ[path_var])
        if not script_path:
            raise Exception('Environment variable %s is set but BNG2.pl could'
                            ' not be found there' % path_var)
    # If the environment variable isn't set, check the standard locations
    else:
        for dist_dir in dist_dirs:
            script_path = check_dist_dir(dist_dir)
            if script_path:
                break
        else:
            raise Exception('Could not find BioNetGen installed in one of the '
                            'following locations:' +
                            ''.join('\n    ' + d for d in dist_dirs))
    # Cache path for future use
    _bng_path = script_path
    return script_path

class GenerateNetworkError(RuntimeError):
    """BNG reported an error when trying to generate a network for a model."""
    pass

_generate_network_code = """
begin actions
generate_network({overwrite=>1});
end actions
"""

def _parse_bng_outfile(out_filename):
    """
    Load and return data from a BNG .gdat or .cdat output file.

    The format of the output files is an initial line of the form::

        #   time   obs1    obs2    obs3  ...

    The column headers are separated by a differing number of spaces (not tabs).
    This function parses out the column names and then uses the numpy.loadtxt
    method to load the outputfile into a numpy record array.

    Parameters
    ----------
    out_filename : string
        Path of file to load.

    """

    try:
        out_file = open(out_filename, 'r')

        line = out_file.readline().strip() # Get the first line
        out_file.close()
        line = line[2:]  # strip off opening '# '
        raw_names = re.split('\s*', line)
        column_names = [raw_name for raw_name in raw_names if not raw_name == '']

        # Create the dtype argument for the numpy record array
        dt = zip(column_names, ('float',)*len(column_names))

        # Load the output file as a numpy record array, skip the name row
        arr = numpy.loadtxt(out_filename, dtype=dt, skiprows=1)

    except Exception as e:
        # FIXME special Exception/Error?
        raise Exception("problem parsing BNG outfile: " + str(e))

    return arr

def run_ssa(model, t_end=10, n_steps=100, output_dir='/tmp', cleanup=True):
    """
    Simulate a model with BNG's SSA simulator and return the trajectories.

    Parameters
    ----------
    model : Model
        Model to simulate.
    t_end : number, optional
        Final time point of the simulation.
    n_steps : int, optional
        Number of steps in the simulation.
    output_dir : string, optional
        Location for temporary files generated by BNG. Defaults to '/tmp'.
    cleanup : bool, optional
        If True (default), delete the temporary files after the simulation is
        finished. If False, leave them in place (in `output_dir`). Useful for
        debugging.

    """

    run_ssa_code = """
    begin actions
    generate_network({overwrite=>1});
    simulate_ssa({t_end=>%f, n_steps=>%d});\n
    end actions
    """ % (t_end, n_steps)

    gen = BngGenerator(model)
    bng_filename = '%s_%d_%d_temp.bngl' % (model.name,
                            os.getpid(), random.randint(0, 10000))
    gdat_filename = bng_filename.replace('.bngl', '.gdat')
    cdat_filename = bng_filename.replace('.bngl', '.cdat')
    net_filename = bng_filename.replace('.bngl', '.net')

    output = StringIO()

    try:
        working_dir = os.getcwd()
        os.chdir(output_dir)
        bng_file = open(bng_filename, 'w')
        bng_file.write(gen.get_content())
        bng_file.write(run_ssa_code)
        bng_file.close()
        p = subprocess.Popen(['perl', _get_bng_path(), bng_filename],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (p_out, p_err) = p.communicate()
        if p.returncode:
            raise GenerateNetworkError(p_out.rstrip("at line")+"\n"+p_err.rstrip())

        output_arr = _parse_bng_outfile(gdat_filename)
        #ssa_file = open(ssa_filename, 'r')
        #output.write(ssa_file.read())
        #net_file.close()
        #if append_stdout:
        #    output.write("#\n# BioNetGen execution log follows\n# ==========")
        #    output.write(re.sub(r'(^|\n)', r'\n# ', p_out))
    finally:
        if cleanup:
            for filename in [bng_filename, gdat_filename,
                             cdat_filename, net_filename]:
                if os.access(filename, os.F_OK):
                    os.unlink(filename)
        os.chdir(working_dir)
    return output_arr

def generate_network(model, cleanup=True, append_stdout=False):
    """
    Return the output from BNG's generate_network function given a model.

    The output is a BNGL model definition with additional sections 'reactions'
    and 'groups', and the 'species' section expanded to contain all possible
    species. BNG refers to this as a 'net' file.

    Parameters
    ----------
    model : Model
        Model to pass to generate_network.
    cleanup : bool, optional
        If True (default), delete the temporary files after the simulation is
        finished. If False, leave them in place (in `output_dir`). Useful for
        debugging.
    append_stdout : bool, optional
        If True, provide BNG.pl's standard output stream as comment lines
        appended to the net file contents. If False (default), do not append it.

    """
    gen = BngGenerator(model)
    bng_filename = '%s_%d_%d_temp.bngl' % (model.name, os.getpid(), random.randint(0, 10000))
    net_filename = bng_filename.replace('.bngl', '.net')
    output = StringIO()
    try:
        bng_file = open(bng_filename, 'w')
        bng_file.write(gen.get_content())
        bng_file.write(_generate_network_code)
        bng_file.close()
        p = subprocess.Popen(['perl', _get_bng_path(), bng_filename],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (p_out, p_err) = p.communicate()
        if p.returncode:
            raise GenerateNetworkError(p_out.rstrip()+"\n"+p_err.rstrip())
        net_file = open(net_filename, 'r')
        output.write(net_file.read())
        net_file.close()
        if append_stdout:
            output.write("#\n# BioNetGen execution log follows\n# ==========")
            output.write(re.sub(r'(^|\n)', r'\n# ', p_out))
    finally:
        if cleanup:
            for filename in [bng_filename, net_filename]:
                if os.access(filename, os.F_OK):
                    os.unlink(filename)
    return output.getvalue()

def generate_equations(model):
    """
    Generate math expressions for reaction rates and species in a model.

    This fills in the following pieces of the model:

    * odes
    * species
    * reactions
    * reactions_bidirectional
    * observables (just `coefficients` and `species` fields for each element)

    """
    # only need to do this once
    # TODO track "dirty" state, i.e. "has model been modified?"
    #   or, use a separate "math model" object to contain ODEs
    if model.odes:
        return
    lines = iter(generate_network(model).split('\n'))
    try:
        while 'begin species' not in lines.next():
            pass
        model.species = []
        while True:
            line = lines.next()
            if 'end species' in line: break
            _parse_species(model, line)

        while 'begin reactions' not in lines.next():
            pass
        model.odes = [sympy.S(0)] * len(model.species)
        reaction_cache = {}
        while True:
            line = lines.next()
            if 'end reactions' in line: break
            (number, reactants, products, rate, rule) = line.strip().split()
            # the -1 is to switch from one-based to zero-based indexing
            reactants = tuple(int(r) - 1 for r in reactants.split(','))
            products = tuple(int(p) - 1 for p in products.split(','))
            rate = rate.rsplit('*')
            (rule_name, is_reverse) = re.match(r'#(\w+)(?:\((reverse)\))?', rule).groups()
            is_reverse = bool(is_reverse)
            r_names = ['s%d' % r for r in reactants]
            combined_rate = sympy.Mul(*[sympy.S(t) for t in r_names + rate])
            rule = model.rules[rule_name]
            reaction = {
                'reactants': reactants,
                'products': products,
                'rate': combined_rate,
                'rule': rule_name,
                'reverse': is_reverse,
                }
            model.reactions.append(reaction)
            key = (rule_name, reactants, products)
            key_reverse = (rule_name, products, reactants)
            reaction_bd = reaction_cache.get(key_reverse)
            if reaction_bd is None:
                # make a copy of the reaction dict
                reaction_bd = dict(reaction)
                # default to false until we find a matching reverse reaction
                reaction_bd['reversible'] = False
                reaction_cache[key] = reaction_bd
                model.reactions_bidirectional.append(reaction_bd)
            else:
                reaction_bd['reversible'] = True
                reaction_bd['rate'] -= combined_rate
            for p in products:
                model.odes[p] += combined_rate
            for r in reactants:
                model.odes[r] -= combined_rate
        # fix up reactions whose reverse version we saw first
        for r in model.reactions_bidirectional:
            if r['reverse']:
                r['reactants'], r['products'] = r['products'], r['reactants']
                r['rate'] *= -1
            # now the 'reverse' value is no longer needed
            del r['reverse']

        while 'begin groups' not in lines.next():
            pass
        while True:
            line = lines.next()
            if 'end groups' in line: break
            _parse_group(model, line)

    except StopIteration as e:
        pass

def _parse_species(model, line):
    """Parse a 'species' line from a BNGL net file."""
    index, species, value = line.strip().split()
    complex_compartment_name, complex_string = re.match(r'(?:@(\w+)::)?(.*)', species).groups()
    monomer_strings = complex_string.split('.')
    monomer_patterns = []
    for ms in monomer_strings:
        monomer_name, site_strings, monomer_compartment_name = re.match(r'(\w+)\(([^)]*)\)(?:@(\w+))?', ms).groups()
        site_conditions = {}
        if len(site_strings):
            for ss in site_strings.split(','):
                # FIXME this should probably be done with regular expressions
                if '!' in ss and '~' in ss:
                    site_name, condition = ss.split('~')
                    state, bond = condition.split('!')
                    if bond == '?':
                        bond = pysb.core.WILD
                    elif bond == '!':
                        bond = pysb.core.ANY
                    else:
                        bond = int(bond)
                    condition = (state, bond)
                elif '!' in ss:
                    site_name, condition = ss.split('!', 1)
                    if '!' in condition:
                        condition = [int(c) for c in condition.split('!')]
                    else:
                        condition = int(condition)
                elif '~' in ss:
                    site_name, condition = ss.split('~')
                else:
                    site_name, condition = ss, None
                site_conditions[site_name] = condition
        monomer = model.monomers[monomer_name]
        monomer_compartment = model.compartments.get(monomer_compartment_name)
        mp = pysb.core.MonomerPattern(monomer, site_conditions, monomer_compartment)
        monomer_patterns.append(mp)

    complex_compartment = model.compartments.get(complex_compartment_name)
    cp = pysb.core.ComplexPattern(monomer_patterns, complex_compartment)
    model.species.append(cp)

def _parse_group(model, line):
    """Parse a 'group' line from a BNGL net file."""
    # values are number (which we ignore), name, and species list
    values = line.strip().split()
    obs = model.observables[values[1]]
    if len(values) == 3:
        # combination is a comma separated list of [coeff*]speciesnumber
        for product in values[2].split(','):
            terms = product.split('*')
            # if no coeff given (just species), insert a coeff of 1
            if len(terms) == 1:
                terms.insert(0, 1)
            obs.coefficients.append(int(terms[0]))
            # -1 to change to 0-based indexing
            obs.species.append(int(terms[1]) - 1)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, unicode_literals
import logging
from .requirements import RequirementsBundle
from .providers.github import Provider as GithubProvider
from .errors import NoPermissionError

logger = logging.getLogger(__name__)

class Bot(object):
    def __init__(self, repo, user_token, bot_token=None,
                 provider=GithubProvider, bundle=RequirementsBundle):
        self.bot_token = bot_token
        self.req_bundle = bundle()
        self.provider = provider(self.req_bundle)
        self.user_token = user_token
        self.bot_token = bot_token
        self.fetched_files = []
        self.repo_name = repo

        self._user = None
        self._user_repo = None
        self._bot = None
        self._bot_repo = None

        self._pull_requests = None

    @property
    def user_repo(self):
        if self._user_repo is None:
            self._user_repo = self.provider.get_repo(token=self.user_token, name=self.repo_name)
        return self._user_repo

    @property
    def user(self):
        if self._user is None:
            self._user = self.provider.get_user(token=self.user_token)
        return self._user

    @property
    def bot(self):
        if self._bot is None:
            self._bot = self.provider.get_user(token=self.bot_token)
        return self._bot

    @property
    def bot_repo(self):
        if self._bot_repo is None:
            self._bot_repo = self.provider.get_repo(token=self.bot_token, name=self.repo_name)
        return self._bot_repo

    @property
    def pull_requests(self):
        if self._pull_requests is None:
            self._pull_requests = [pr for pr in self.provider.iter_issues(repo=self.user_repo,
                                                                          creator=self.bot)]
        return self._pull_requests

    def update(self, branch=None, initial=True, pin_unpinned=False):

        if branch is None:
            branch = self.provider.get_default_branch(repo=self.user_repo)

        self.get_all_requirements(branch=branch)

        #
        self.apply_updates(branch, initial=initial, pin_unpinned=pin_unpinned)

        return self.req_bundle

    def apply_updates(self, branch, initial, pin_unpinned):

        for title, body, update_branch, updates in \
                self.req_bundle.get_updates(inital=initial, pin_unpinned=pin_unpinned):
            if title not in [pr.title for pr in self.pull_requests]:
                pull_request = self.commit_and_pull(
                    base_branch=branch,
                    new_branch=update_branch,
                    title=title,
                    body=body,
                    updates=updates
                )
            else:
                pull_request = next((pr for pr in self.pull_requests if pr.title == title), None)

            for update in updates:
                update.requirement.pull_request = pull_request

    def commit_and_pull(self, base_branch, new_branch, title, body, updates):

        # create new branch
        self.provider.create_branch(
            base_branch=base_branch,
            new_branch=new_branch,
            repo=self.user_repo
        )

        updated_files = {}
        for update in updates:

            if update.requirement_file.path in updated_files:
                sha = updated_files[update.requirement_file.path]["sha"]
                content = updated_files[update.requirement_file.path]["content"]
            else:
                sha = update.requirement_file.sha
                content = update.requirement_file.content

            content = update.requirement.update_content(content)
            new_sha = self.provider.create_commit(
                repo=self.user_repo,
                path=update.requirement_file.path,
                branch=new_branch,
                content=content,
                commit_message=update.commit_message,
                sha=sha,
                committer=self.bot if self.bot_token else self.user,
            )

            updated_files[update.requirement_file.path] = {"sha": new_sha, "content": content}

        return self.create_pull_request(
            title=title,
            body=body,
            base_branch=base_branch,
            new_branch=new_branch
        )

    def create_pull_request(self, title, body, base_branch, new_branch):

        # if we have a bot user that creates the PR, we might run into problems on private
        # repos because the bot has to be a collaborator. We try to submit the PR before checking
        # the permissions because that saves us API calls in most cases
        if self.bot_token:
            try:
                return self.provider.create_pull_request(
                    repo=self.bot_repo,
                    title=title,
                    body=body,
                    base_branch=base_branch,
                    new_branch=new_branch,
                )
            except NoPermissionError:
                self.provider.get_pull_request_permissions(self.bot, self.user_repo)

        return self.provider.create_pull_request(
            repo=self.bot_repo if self.bot_token else self.user_repo,
            title=title,
            body=body,
            base_branch=base_branch,
            new_branch=new_branch,
        )

    def get_all_requirements(self, branch):
        for file_type, path in self.provider.iter_git_tree(branch=branch, repo=self.user_repo):
            if file_type == "blob":
                if "requirements" in path:
                    if path.endswith("txt") or path.endswith("pip"):
                        self.add_requirement_file(path)

    def add_requirement_file(self, path):
        if not self.req_bundle.has_file_in_path(path):
            req_file = self.provider.get_requirement_file(path=path, repo=self.user_repo)
            if req_file is not None:
                self.req_bundle.append(req_file)
                for other_file in req_file.other_files:
                    self.add_requirement_file(other_file)

class DryBot(Bot):
    def commit_and_pull(self, base_branch, new_branch, title, body, updates):  # pragma: no cover
        return None
<commit_msg>fixed typo<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, unicode_literals
import logging
from .requirements import RequirementsBundle
from .providers.github import Provider as GithubProvider
from .errors import NoPermissionError

logger = logging.getLogger(__name__)

class Bot(object):
    def __init__(self, repo, user_token, bot_token=None,
                 provider=GithubProvider, bundle=RequirementsBundle):
        self.bot_token = bot_token
        self.req_bundle = bundle()
        self.provider = provider(self.req_bundle)
        self.user_token = user_token
        self.bot_token = bot_token
        self.fetched_files = []
        self.repo_name = repo

        self._user = None
        self._user_repo = None
        self._bot = None
        self._bot_repo = None

        self._pull_requests = None

    @property
    def user_repo(self):
        if self._user_repo is None:
            self._user_repo = self.provider.get_repo(token=self.user_token, name=self.repo_name)
        return self._user_repo

    @property
    def user(self):
        if self._user is None:
            self._user = self.provider.get_user(token=self.user_token)
        return self._user

    @property
    def bot(self):
        if self._bot is None:
            self._bot = self.provider.get_user(token=self.bot_token)
        return self._bot

    @property
    def bot_repo(self):
        if self._bot_repo is None:
            self._bot_repo = self.provider.get_repo(token=self.bot_token, name=self.repo_name)
        return self._bot_repo

    @property
    def pull_requests(self):
        if self._pull_requests is None:
            self._pull_requests = [pr for pr in self.provider.iter_issues(repo=self.user_repo,
                                                                          creator=self.bot)]
        return self._pull_requests

    def update(self, branch=None, initial=True, pin_unpinned=False):

        if branch is None:
            branch = self.provider.get_default_branch(repo=self.user_repo)

        self.get_all_requirements(branch=branch)

        #
        self.apply_updates(branch, initial=initial, pin_unpinned=pin_unpinned)

        return self.req_bundle

    def apply_updates(self, branch, initial, pin_unpinned):
        # check if we have an initial PR open. If this is the case, we attach the initial PR
        # to all updates and are done. The `Initial Update` has to be merged (or at least closed)
        # before we continue to do anything here.
        initial_pr = next(
            (pr for pr in self.pull_requests if pr.title == "Initial Update" and pr.is_open),
            False
        )

        for title, body, update_branch, updates in \
                self.req_bundle.get_updates(initial=initial, pin_unpinned=pin_unpinned):
            if initial_pr:
                pull_request = initial_pr
            elif title not in [pr.title for pr in self.pull_requests]:
                pull_request = self.commit_and_pull(
                    base_branch=branch,
                    new_branch=update_branch,
                    title=title,
                    body=body,
                    updates=updates
                )
            else:
                pull_request = next((pr for pr in self.pull_requests if pr.title == title), None)

            for update in updates:
                update.requirement.pull_request = pull_request

    def commit_and_pull(self, base_branch, new_branch, title, body, updates):

        # create new branch
        self.provider.create_branch(
            base_branch=base_branch,
            new_branch=new_branch,
            repo=self.user_repo
        )

        updated_files = {}
        for update in updates:

            if update.requirement_file.path in updated_files:
                sha = updated_files[update.requirement_file.path]["sha"]
                content = updated_files[update.requirement_file.path]["content"]
            else:
                sha = update.requirement_file.sha
                content = update.requirement_file.content

            content = update.requirement.update_content(content)
            new_sha = self.provider.create_commit(
                repo=self.user_repo,
                path=update.requirement_file.path,
                branch=new_branch,
                content=content,
                commit_message=update.commit_message,
                sha=sha,
                committer=self.bot if self.bot_token else self.user,
            )

            updated_files[update.requirement_file.path] = {"sha": new_sha, "content": content}

        return self.create_pull_request(
            title=title,
            body=body,
            base_branch=base_branch,
            new_branch=new_branch
        )

    def create_pull_request(self, title, body, base_branch, new_branch):

        # if we have a bot user that creates the PR, we might run into problems on private
        # repos because the bot has to be a collaborator. We try to submit the PR before checking
        # the permissions because that saves us API calls in most cases
        if self.bot_token:
            try:
                return self.provider.create_pull_request(
                    repo=self.bot_repo,
                    title=title,
                    body=body,
                    base_branch=base_branch,
                    new_branch=new_branch,
                )
            except NoPermissionError:
                self.provider.get_pull_request_permissions(self.bot, self.user_repo)

        return self.provider.create_pull_request(
            repo=self.bot_repo if self.bot_token else self.user_repo,
            title=title,
            body=body,
            base_branch=base_branch,
            new_branch=new_branch,
        )

    def get_all_requirements(self, branch):
        for file_type, path in self.provider.iter_git_tree(branch=branch, repo=self.user_repo):
            if file_type == "blob":
                if "requirements" in path:
                    if path.endswith("txt") or path.endswith("pip"):
                        self.add_requirement_file(path)

    def add_requirement_file(self, path):
        if not self.req_bundle.has_file_in_path(path):
            req_file = self.provider.get_requirement_file(path=path, repo=self.user_repo)
            if req_file is not None:
                self.req_bundle.append(req_file)
                for other_file in req_file.other_files:
                    self.add_requirement_file(other_file)

class DryBot(Bot):
    def commit_and_pull(self, base_branch, new_branch, title, body, updates):  # pragma: no cover
        return None
<|endoftext|>