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

# Copyright (c) 2015, Palo Alto Networks
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

# Author: Brian Torres-Gil <btorres-gil@paloaltonetworks.com>

"""Device updater handles software versions and updates for devices

About this script
-----------------
Given an IP address, tags or identifies the ip address in PAN-OS.
Behavior in PAN-OS 6.0 and higher:
  IP address is tagged. IP's can have multiple tags. Tags are
  referenced in Dynamic Address Groups. Boolean logic (AND/OR)
  can be used in Dynamic Address Groups to filter ip's based on tags
Behavior in PAN-OS 5.1 and lower:
  IP address is mapped to an identifier. Only one identifier
  per IP address. The identifier associates the IP with a single
  Dynamic Address Object in PAN-OS.

The script assumes that you have firewall policy setup that blocks
traffic for a given dynamic address group

It is important to recognize that this script does NOT modify a firewall rule!
It is only tagging IP addresses with metadata so that the existing firewall
policy can act accordingly.

How to use this script
----------------------
In the example below, we are blocking all ip's returned by the search
example1: index=pan_logs 1.1.1.1 | stats dc(dst_ip) by dst_ip | pantag action="add" tag="malware-infected" device="1.0.0.1"
Adds a 'malware-infected' tag to the IP 1.1.1.1 on the firewall with ip 1.0.0.1
example2: index=pan_logs wine | stats dc(dst_ip) by dst_ip | pantag action="rem" group="shairpoint" device="sales-fw"
Removes the 'shairpoint' tag from all dst_ip returned by the search on the firewall with hostname sales-fw

"""

#########################################################
# Do NOT modify anything below this line unless you are
# certain of the ramifications of the changes
#########################################################

import sys  # for system params and sys.exit()
import os

libpath = os.path.dirname(os.path.abspath(__file__))
sys.path[:0] = [os.path.join(libpath, 'lib')]
import common
import environment

logger = common.logging.getLogger().getChild('panTag')
#logger.setLevel(pan_common.logging.INFO)

try:
    if environment.run_by_splunk():
        import splunk.Intersplunk  # so you can interact with Splunk
        import splunk.entity as entity  # for splunk config info

    libpath = os.path.dirname(os.path.abspath(__file__))
    sys.path[:0] = [os.path.join(libpath, 'lib')]
    sys.path[:0] = [os.path.join(libpath, 'lib', 'pan-python', 'lib')]
    sys.path[:0] = [os.path.join(libpath, 'lib', 'pandevice')]
    import pandevice
    from pandevice.panorama import Panorama
    from pandevice.firewall import Firewall
    import pan.xapi

    from common import log

except Exception as e:
    # Handle exception to produce logs to python.log
    common.exit_with_error(e)

# Default fields that contain IP addresses and should be tagged if they exist
IP_FIELDS = ['src_ip', 'dest_ip', 'ip']

def main_splunk():
    # Get arguments
    args, kwargs = splunk.Intersplunk.getKeywordsAndOptions()

    # Enable debugging by passing 'debug=yes' as an argument of
    # the command on the Splunk searchbar.

    debug = common.check_debug(kwargs)

    # kwargs contains important parameters.
    # parameters from splunk searchbar include:
    #   action
    #   device
    #   panorama
    #   serial
    #   vsys
    #   tag
    #   tag_field
    #   ip_field
    #   debug

    # Verify required args were passed to command
    log(debug, "Determining if required arguments are present")
    if 'device' not in kwargs and 'panorama' not in kwargs:
        common.exit_with_error("pantag: Missing required command argument: device or panorama", 3)
    if 'tag' not in kwargs and 'tag_field' not in kwargs:
        common.exit_with_error("pantag: Missing required command argument: tag or tag_field", 3)

    # Assign defaults to fields that aren't specified
    action = kwargs['action'] if 'action' in kwargs else "add"
    vsys = kwargs['vsys'] if 'vsys' in kwargs else None
    ip_field = kwargs['ip_field'] if 'ip_field' in kwargs else "src_ip"
    tag = kwargs['tag'] if 'tag' in kwargs else None
    tag_field = kwargs['tag_field'] if 'tag_field' in kwargs else None

    # Determine if device hostname or serial was provided as argument or should be pulled from entries
    log(debug, "Determining how firewalls should be contacted based on arguments")
    use_panorama = False
    hostname = None
    serial = None
    if "device" in kwargs:
        hostname = kwargs['device']
        if vsys is None:
            vsys = "vsys1"
    elif "panorama" in kwargs:
        use_panorama = True
        hostname = kwargs['panorama']
        if "serial" in kwargs:
            serial = kwargs['serial']
            if vsys is None:
                vsys = "vsys1"
    else:
        common.exit_with_error("pantag: Missing required command argument: device or panorama", 3)
    log(debug, "Use Panorama: %s" % use_panorama)
    log(debug, "VSys: %s" % vsys)
    log(debug, "Hostname: %s" % hostname)
    if use_panorama and serial is not None:
        log(debug, "Device Serial: %s" % serial)
    else:
        log(debug, "Using serials from logs")

    # Results contains the data from the search results and settings
    # contains the sessionKey that we can use to talk to Splunk
    results, unused1, settings = splunk.Intersplunk.getOrganizedResults()
    # Get the sessionKey
    sessionKey = settings['sessionKey']

    log(debug, "Begin get API key")
    # Get the API key from the Splunk store or from the device at hostname if no apikey is stored
    apikey = common.apikey(sessionKey, hostname, debug)

    # Create the connection to the firewall or Panorama
    panorama=None
    if use_panorama:
        # For Panorama, create the Panorama object, and the firewall if only one serial
        panorama = Panorama(hostname, api_key=apikey)
        if serial is not None:
            firewall = {'firewall': Firewall(panorama=panorama, serial=serial, vsys=vsys)}
            firewall['firewall'].userid.batch_start()
        else:
            firewall = {}
    else:
        firewall = {'firewall': Firewall(hostname, api_key=apikey, vsys=vsys)}
        firewall['firewall'].userid.batch_start()

    # Collect all the ip addresses and tags into firewall batch requests
    for result in results:

        ## Find the serial (if not a single firewall)

        if use_panorama and serial is None:
            try:
                this_serial = result['serial_number']
                this_vsys = result['vsys']
            except KeyError as e:
                result['status'] = "ERROR: Unable to determine serial number or vsys of device"
                continue
            else:
                ## Create the firewall object if using multiple serials
                if this_serial in firewall:
                    this_firewall = firewall[(this_serial, this_vsys)]
                else:
                    # Create the firewall object for this serial
                    firewall[(this_serial, this_vsys)] = Firewall(panorama=panorama, serial=this_serial, vsys=this_vsys)
                    this_firewall = firewall[(this_serial, this_vsys)]
                    this_firewall.userid.batch_start()
        else:
            this_firewall = firewall

        ## Find the tag (if a tag_field was specified)

        this_tag = []
        if tag_field is not None:
            try:
                this_tag.append(result[tag_field])
            except KeyError as e:
                result['status'] = "ERROR: Unable to determine tag from field: %s" % tag_field
                continue
        if tag is not None:
            this_tag.append(tag)

        ## Find the IP

        try:
            this_ip = result[ip_field]
        except KeyError as e:
            result['status'] = "ERROR: Unable to determine ip from field: %s" % ip_field

        ## Create a request in the batch user-id update for the firewall
        ## No API call to the firewall happens until all batch requests are created.

        if action == "add":
            log(debug, "Registering tags on firewall %s: %s - %s" % (this_firewall, this_ip, this_tag))
            this_firewall.userid.register(this_ip, this_tag)
        else:
            log(debug, "Unregistering tags on firewall %s: %s - %s" % (this_firewall, this_ip, this_tag))
            this_firewall.userid.unregister(this_ip, this_tag)

        result['status'] = "Submitted successfully"

    ## Make the API calls to the User-ID API of each firewall

    try:
        for fw in firewall.values():
            fw.userid.batch_end()

    except pan.xapi.PanXapiError as e:
        common.exit_with_error(str(e))

    except Exception as e:
        common.exit_with_error(str(e))

    # output results
    splunk.Intersplunk.outputResults(results)

def main_cli():
    pass

if __name__ == "__main__":
    if environment.run_by_splunk():
        main_splunk()
    else:
        main_cli()
<commit_msg>Minor cleanup to pantag script<commit_after>#!/usr/bin/env python

# Copyright (c) 2015, Palo Alto Networks
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

# Author: Brian Torres-Gil <btorres-gil@paloaltonetworks.com>

"""Device updater handles software versions and updates for devices

About this script
-----------------
Given an IP address, tags or identifies the ip address in PAN-OS.
Behavior in PAN-OS 6.0 and higher:
  IP address is tagged. IP's can have multiple tags. Tags are
  referenced in Dynamic Address Groups. Boolean logic (AND/OR)
  can be used in Dynamic Address Groups to filter ip's based on tags
Behavior in PAN-OS 5.1 and lower:
  IP address is mapped to an identifier. Only one identifier
  per IP address. The identifier associates the IP with a single
  Dynamic Address Object in PAN-OS.

The script assumes that you have firewall policy setup that blocks
traffic for a given dynamic address group

It is important to recognize that this script does NOT modify a firewall rule!
It is only tagging IP addresses with metadata so that the existing firewall
policy can act accordingly.

This script supports connection to a firewall or to Panorama.
"""

#########################################################
# Do NOT modify anything below this line unless you are
# certain of the ramifications of the changes
#########################################################

import sys  # for system params and sys.exit()
import os

libpath = os.path.dirname(os.path.abspath(__file__))
sys.path[:0] = [os.path.join(libpath, 'lib')]
import common
import environment

logger = common.logging.getLogger().getChild('panTag')
#logger.setLevel(pan_common.logging.INFO)

try:
    if environment.run_by_splunk():
        import splunk.Intersplunk  # so you can interact with Splunk
        import splunk.entity as entity  # for splunk config info

    libpath = os.path.dirname(os.path.abspath(__file__))
    sys.path[:0] = [os.path.join(libpath, 'lib')]
    sys.path[:0] = [os.path.join(libpath, 'lib', 'pan-python', 'lib')]
    sys.path[:0] = [os.path.join(libpath, 'lib', 'pandevice')]
    import pandevice
    from pandevice.panorama import Panorama
    from pandevice.firewall import Firewall
    import pan.xapi

    from common import log

except Exception as e:
    # Handle exception to produce logs to python.log
    common.exit_with_error(e)

# Default fields that contain IP addresses and should be tagged if they exist
IP_FIELDS = ['src_ip', 'dest_ip', 'ip']

def main_splunk():
    # Get arguments
    args, kwargs = splunk.Intersplunk.getKeywordsAndOptions()

    # Enable debugging by passing 'debug=yes' as an argument of
    # the command on the Splunk searchbar.

    debug = common.check_debug(kwargs)

    # kwargs contains important parameters.
    # parameters from splunk searchbar include:
    #   action
    #   device
    #   panorama
    #   serial
    #   vsys
    #   tag
    #   tag_field
    #   ip_field
    #   debug

    # Verify required args were passed to command
    log(debug, "Determining if required arguments are present")
    if 'device' not in kwargs and 'panorama' not in kwargs:
        common.exit_with_error("Missing required command argument: device or panorama", 3)
    if 'tag' not in kwargs and 'tag_field' not in kwargs:
        common.exit_with_error("Missing required command argument: tag or tag_field", 3)

    # Assign defaults to fields that aren't specified
    action = kwargs['action'] if 'action' in kwargs else "add"
    vsys = kwargs['vsys'] if 'vsys' in kwargs else None
    ip_field = kwargs['ip_field'] if 'ip_field' in kwargs else "src_ip"
    tag = kwargs['tag'] if 'tag' in kwargs else None
    tag_field = kwargs['tag_field'] if 'tag_field' in kwargs else None

    # Determine if device hostname or serial was provided as argument or should be pulled from entries
    log(debug, "Determining how firewalls should be contacted based on arguments")
    use_panorama = False
    hostname = None
    serial = None
    if "device" in kwargs:
        hostname = kwargs['device']
        if vsys is None:
            vsys = "vsys1"
    elif "panorama" in kwargs:
        use_panorama = True
        hostname = kwargs['panorama']
        if "serial" in kwargs:
            serial = kwargs['serial']
            if vsys is None:
                vsys = "vsys1"
    else:
        common.exit_with_error("Missing required command argument: device or panorama", 3)
    log(debug, "Use Panorama: %s" % use_panorama)
    log(debug, "VSys: %s" % vsys)
    log(debug, "Hostname: %s" % hostname)
    if use_panorama and serial is not None:
        log(debug, "Device Serial: %s" % serial)
    else:
        log(debug, "Using serials from logs")

    # Results contains the data from the search results and settings
    # contains the sessionKey that we can use to talk to Splunk
    results, unused1, settings = splunk.Intersplunk.getOrganizedResults()
    # Get the sessionKey
    sessionKey = settings['sessionKey']

    log(debug, "Begin get API key")
    # Get the API key from the Splunk store or from the device at hostname if no apikey is stored
    apikey = common.apikey(sessionKey, hostname, debug)

    # Create the connection to the firewall or Panorama
    panorama = None
    if use_panorama:
        # For Panorama, create the Panorama object, and the firewall if only one serial
        panorama = Panorama(hostname, api_key=apikey)
        if serial is not None:
            firewall = {'firewall': Firewall(panorama=panorama, serial=serial, vsys=vsys)}
            firewall['firewall'].userid.batch_start()
        else:
            firewall = {}
    else:
        firewall = {'firewall': Firewall(hostname, api_key=apikey, vsys=vsys)}
        firewall['firewall'].userid.batch_start()

    # Collect all the ip addresses and tags into firewall batch requests
    for result in results:

        ## Find the serial (if not a single firewall)

        if use_panorama and serial is None:
            try:
                this_serial = result['serial_number']
                this_vsys = result['vsys']
            except KeyError as e:
                result['status'] = "ERROR: Unable to determine serial number or vsys of device"
                continue
            else:
                ## Create the firewall object if using multiple serials
                if this_serial in firewall:
                    this_firewall = firewall[(this_serial, this_vsys)]
                else:
                    # Create the firewall object for this serial
                    firewall[(this_serial, this_vsys)] = Firewall(panorama=panorama, serial=this_serial, vsys=this_vsys)
                    this_firewall = firewall[(this_serial, this_vsys)]
                    this_firewall.userid.batch_start()
        else:
            this_firewall = firewall

        ## Find the tag (if a tag_field was specified)

        this_tag = []
        if tag_field is not None:
            try:
                this_tag.append(result[tag_field])
            except KeyError as e:
                result['status'] = "ERROR: Unable to determine tag from field: %s" % tag_field
                continue
        if tag is not None:
            this_tag.append(tag)

        ## Find the IP

        try:
            this_ip = result[ip_field]
        except KeyError as e:
            result['status'] = "ERROR: Unable to determine ip from field: %s" % ip_field

        ## Create a request in the batch user-id update for the firewall
        ## No API call to the firewall happens until all batch requests are created.

        if action == "add":
            log(debug, "Registering tags on firewall %s: %s - %s" % (this_firewall, this_ip, this_tag))
            this_firewall.userid.register(this_ip, this_tag)
        else:
            log(debug, "Unregistering tags on firewall %s: %s - %s" % (this_firewall, this_ip, this_tag))
            this_firewall.userid.unregister(this_ip, this_tag)

        result['status'] = "Submitted successfully"

    ## Make the API calls to the User-ID API of each firewall

    try:
        for fw in firewall.values():
            fw.userid.batch_end()

    except pan.xapi.PanXapiError as e:
        common.exit_with_error(str(e))

    except Exception as e:
        common.exit_with_error(str(e))

    # output results
    splunk.Intersplunk.outputResults(results)

def main_cli():
    raise NotImplementedError

if __name__ == "__main__":
    if environment.run_by_splunk():
        main_splunk()
    else:
        main_cli()
<|endoftext|>


<commit_before>import codecs
import os
import sys
import time
import re

wpull_hook = globals().get('wpull_hook')  # silence code checkers
tries = 0
total_tries = 0
max_gallery_page = None
max_scraps_page = None
max_favorites_page = None
item_dir = os.environ['item_dir']

def print_(*args, **kwargs):
    print(*args, **kwargs)
    sys.stdout.flush()

def engine_run():
    with open(os.path.join(item_dir, 'usernames.txt'), 'a'):
        pass

def accept_url(url_info, record_info, verdict, reasons):
    global max_gallery_page
    global max_scraps_page

    url = url_info['url']

    if verdict:
        match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

        if match:
            what_type = match.group(1).lower()

            if what_type in ('gallery', 'scraps', 'favorites'):
                num = int(match.group(3))

                if what_type == 'gallery':
                    if max_gallery_page is not None and num > max_gallery_page:
                        print_('Pagination complete for gallery')
                        return False
                elif what_type == 'scraps':
                    if max_scraps_page is not None and num > max_scraps_page:
                        print_('Pagination complete for scraps')
                        return False
                elif what_type == 'favorites':
                    if max_favorites_page is not None and num > max_favorites_page:
                        print_('Pagination complete for favorites')
                        return False
                else:
                    raise Exception('Unknown what type!')

        if '/themes/classic/' in url:
            return False
        if 'google-analytics.com' in url:
            return False
        if 'google.com/analytics' in url:
            return False
        if 'gstatic.com/analytics' in url:
            return False
        if re.search(r'/commissions/.*/(add|manage)/$', url_info['url']):
            return False

        if 'facdn.net' in url and 'furaffinity.net/favorites/' in record_info['referrer']:
            return False

        if 'puu.sh/' in url:
            # ArchiveTeam is banned from stupid puush
            return False

        if 'googleadservices.com' in url:
            return False
        if 'googlesyndication.com/pagead/' in url:
            return False
        if 'googletagservices.com' in url:
            return False

    if not verdict and 'facdn.net' in url and 'furaffinity.net/view/' in record_info['referrer']:
        return True

    return verdict

def handle_pre_response(url_info, record_info, response_info):
    return wpull_hook.actions.NORMAL

def handle_response(url_info, record_info, response_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 5000:
        raise Exception('Too many tries in this session!')

    status_code = response_info['status_code']
    url = url_info['url']

    if status_code != 404 and status_code >= 400:
        if 'furaffinity.net' in url or 'facdn.net' in url:
            print_('Uh oh!. Sleeping...')
            time.sleep(60)

        tries += 1

        if tries >= 5 and ('furaffinity.net' in url or 'facdn.net' in url):
            raise Exception('Giving up')
        elif tries >= 5:
            return wpull_hook.actions.FINISH
        else:
            return wpull_hook.actions.RETRY

    tries = 0

    return wpull_hook.actions.NORMAL

def handle_error(url_info, record_info, error_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 5000:
        raise Exception('Too many tries in this session!')

    tries += 1
    url = url_info['url']
    if tries >= 5 and ('furaffinity.net' in url or 'facdn.net' in url):
        raise Exception('Giving up')
    elif tries >= 5:
        return wpull_hook.actions.FINISH

    return wpull_hook.actions.NORMAL

def get_urls(filename, url_info, document_info):
    urls = []

    if 'furaffinity.net' in url_info['hostname']:
        with open(filename, 'r', encoding='utf8', errors='replace') as file:
            text = file.read(1048576)

        check_ok_content(text)

        if not is_text_404(text):
            with open(os.path.join(item_dir, 'usernames.txt'), 'a', encoding='utf8', errors='replace') as file:
                for username in scrape_usernames(text):
                    file.write(username)
                    file.write('\n')

            url = url_info['url']
            check_pagination(text, url)

            if 'furaffinity.net/view/' in url:
                check_full_view(text)

        print_('Looking good so far..')

        match = re.search(r'Download', text)

        if match:
            urls.append({
                'url': match.group(1)
            })

    return urls

def check_ok_content(text):
    ok_text_found = (
        'Commission Info' in text or
        'Submission information:' in text or
        'This user cannot be found.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )

    is_404_error_page = is_text_404(text)

    if ok_text_found and not is_404_error_page and '/logout/' not in text:
        print_('Problem detected. Not logged in! Sleeping.')
        time.sleep(60)
        raise Exception('Not logged in!')
    elif ok_text_found and not is_404_error_page and 'Toggle to hide Mature and Adult submissions.' not in text:
        print_('Problem detected. Cannot view adult material! Sleeping.')
        time.sleep(60)
        raise Exception('Cannot view adult material!')

def check_full_view(text):
    if ' type="application/x-shockwave-flash"' in text and \
            'facdn.net/art/' in text and \
            ('">Download' in text or '">Download ' in text):
        return

    if 'var is_full = 1;' not in text:
        raise Exception('Full view not found!')

def is_text_404(text):
    is_404_error_page = (
        'This user cannot be found.' in text or
        'This user has voluntarily disabled access to their userpage.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )
    return is_404_error_page

def scrape_usernames(text):
    for match in re.finditer(r'href="/user/([^"]+)"', text):
        username = match.group(1)
        username = username.strip('/')
        yield username

def check_pagination(text, url):
    global max_gallery_page
    global max_scraps_page
    global max_favorites_page

    match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

    if match:
        # if codecs.encode('kzyuggc_znxrf_gebtqbe_fnq', 'rot_13') not in text:
        #     raise Exception('Could not find pagination form!')

        what_type = match.group(1).lower()

        if what_type in ('gallery', 'scraps', 'favorites') and 'There are no submissions to list' in text:
            num = int(match.group(3))

            if what_type == 'gallery':
                if max_gallery_page is None:
                    max_gallery_page = num
            elif what_type == 'scraps':
                if max_scraps_page is None:
                    max_scraps_page = num
            elif what_type == 'favorites':
                if max_favorites_page is None:
                    max_favorites_page = num

                    if max_favorites_page > 1000:
                        max_favorites_page = 1000
                        print_('Wow, this user likes a lot of things! Capped the pagination.')
            else:
                raise Exception('Unknown what type!')

def wait_time(seconds, url_info, url_record, response, error):
    if 'facdn.net' in url_info['hostname'] and not error:
        return 0
    else:
        return seconds

wpull_hook.callbacks.engine_run = engine_run
# wpull_hook.callbacks.resolve_dns = resolve_dns
wpull_hook.callbacks.accept_url = accept_url
# wpull_hook.callbacks.queued_url = queued_url
# wpull_hook.callbacks.dequeued_url = dequeued_url
wpull_hook.callbacks.handle_pre_response = handle_pre_response
wpull_hook.callbacks.handle_response = handle_response
wpull_hook.callbacks.handle_error = handle_error
wpull_hook.callbacks.get_urls = get_urls
wpull_hook.callbacks.wait_time = wait_time
# wpull_hook.callbacks.finish_statistics = finish_statistics
# wpull_hook.callbacks.exit_status = exit_status
wpull_hook.callbacks.version = 3
<commit_msg>Fix max fav page limit logic.<commit_after>import codecs
import os
import sys
import time
import re

wpull_hook = globals().get('wpull_hook')  # silence code checkers
tries = 0
total_tries = 0
max_gallery_page = None
max_scraps_page = None
max_favorites_page = None
item_dir = os.environ['item_dir']

def print_(*args, **kwargs):
    print(*args, **kwargs)
    sys.stdout.flush()

def engine_run():
    with open(os.path.join(item_dir, 'usernames.txt'), 'a'):
        pass

def accept_url(url_info, record_info, verdict, reasons):
    global max_gallery_page
    global max_scraps_page

    url = url_info['url']

    if verdict:
        match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

        if match:
            what_type = match.group(1).lower()

            if what_type in ('gallery', 'scraps', 'favorites'):
                num = int(match.group(3))

                if what_type == 'gallery':
                    if max_gallery_page is not None and num > max_gallery_page:
                        print_('Pagination complete for gallery')
                        return False
                elif what_type == 'scraps':
                    if max_scraps_page is not None and num > max_scraps_page:
                        print_('Pagination complete for scraps')
                        return False
                elif what_type == 'favorites':
                    if max_favorites_page is not None and num > max_favorites_page:
                        print_('Pagination complete for favorites')
                        return False
                    elif num > 1000:
                        print_('Wow, this user likes a lot of things! Capped the pagination.')
                        return False
                else:
                    raise Exception('Unknown what type!')

        if '/themes/classic/' in url:
            return False
        if 'google-analytics.com' in url:
            return False
        if 'google.com/analytics' in url:
            return False
        if 'gstatic.com/analytics' in url:
            return False
        if re.search(r'/commissions/.*/(add|manage)/$', url_info['url']):
            return False

        if 'facdn.net' in url and 'furaffinity.net/favorites/' in record_info['referrer']:
            return False

        if 'puu.sh/' in url:
            # ArchiveTeam is banned from stupid puush
            return False

        if 'googleadservices.com' in url:
            return False
        if 'googlesyndication.com/pagead/' in url:
            return False
        if 'googletagservices.com' in url:
            return False

    if not verdict and 'facdn.net' in url and 'furaffinity.net/view/' in record_info['referrer']:
        return True

    return verdict

def handle_pre_response(url_info, record_info, response_info):
    return wpull_hook.actions.NORMAL

def handle_response(url_info, record_info, response_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 5000:
        raise Exception('Too many tries in this session!')

    status_code = response_info['status_code']
    url = url_info['url']

    if status_code != 404 and status_code >= 400:
        if 'furaffinity.net' in url or 'facdn.net' in url:
            print_('Uh oh!. Sleeping...')
            time.sleep(60)

        tries += 1

        if tries >= 5 and ('furaffinity.net' in url or 'facdn.net' in url):
            raise Exception('Giving up')
        elif tries >= 5:
            return wpull_hook.actions.FINISH
        else:
            return wpull_hook.actions.RETRY

    tries = 0

    return wpull_hook.actions.NORMAL

def handle_error(url_info, record_info, error_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 5000:
        raise Exception('Too many tries in this session!')

    tries += 1
    url = url_info['url']
    if tries >= 5 and ('furaffinity.net' in url or 'facdn.net' in url):
        raise Exception('Giving up')
    elif tries >= 5:
        return wpull_hook.actions.FINISH

    return wpull_hook.actions.NORMAL

def get_urls(filename, url_info, document_info):
    urls = []

    if 'furaffinity.net' in url_info['hostname']:
        with open(filename, 'r', encoding='utf8', errors='replace') as file:
            text = file.read(1048576)

        check_ok_content(text)

        if not is_text_404(text):
            with open(os.path.join(item_dir, 'usernames.txt'), 'a', encoding='utf8', errors='replace') as file:
                for username in scrape_usernames(text):
                    file.write(username)
                    file.write('\n')

            url = url_info['url']
            check_pagination(text, url)

            if 'furaffinity.net/view/' in url:
                check_full_view(text)

        print_('Looking good so far..')

        match = re.search(r'Download', text)

        if match:
            urls.append({
                'url': match.group(1)
            })

    return urls

def check_ok_content(text):
    ok_text_found = (
        'Commission Info' in text or
        'Submission information:' in text or
        'This user cannot be found.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )

    is_404_error_page = is_text_404(text)

    if ok_text_found and not is_404_error_page and '/logout/' not in text:
        print_('Problem detected. Not logged in! Sleeping.')
        time.sleep(60)
        raise Exception('Not logged in!')
    elif ok_text_found and not is_404_error_page and 'Toggle to hide Mature and Adult submissions.' not in text:
        print_('Problem detected. Cannot view adult material! Sleeping.')
        time.sleep(60)
        raise Exception('Cannot view adult material!')

def check_full_view(text):
    if ' type="application/x-shockwave-flash"' in text and \
            'facdn.net/art/' in text and \
            ('">Download' in text or '">Download ' in text):
        return

    if 'var is_full = 1;' not in text:
        raise Exception('Full view not found!')

def is_text_404(text):
    is_404_error_page = (
        'This user cannot be found.' in text or
        'This user has voluntarily disabled access to their userpage.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )
    return is_404_error_page

def scrape_usernames(text):
    for match in re.finditer(r'href="/user/([^"]+)"', text):
        username = match.group(1)
        username = username.strip('/')
        yield username

def check_pagination(text, url):
    global max_gallery_page
    global max_scraps_page
    global max_favorites_page

    match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

    if match:
        # if codecs.encode('kzyuggc_znxrf_gebtqbe_fnq', 'rot_13') not in text:
        #     raise Exception('Could not find pagination form!')

        what_type = match.group(1).lower()

        if what_type in ('gallery', 'scraps', 'favorites') and 'There are no submissions to list' in text:
            num = int(match.group(3))

            if what_type == 'gallery':
                if max_gallery_page is None:
                    max_gallery_page = num
            elif what_type == 'scraps':
                if max_scraps_page is None:
                    max_scraps_page = num
            elif what_type == 'favorites':
                if max_favorites_page is None:
                    max_favorites_page = num
            else:
                raise Exception('Unknown what type!')

def wait_time(seconds, url_info, url_record, response, error):
    if 'facdn.net' in url_info['hostname'] and not error:
        return 0
    else:
        return seconds

wpull_hook.callbacks.engine_run = engine_run
# wpull_hook.callbacks.resolve_dns = resolve_dns
wpull_hook.callbacks.accept_url = accept_url
# wpull_hook.callbacks.queued_url = queued_url
# wpull_hook.callbacks.dequeued_url = dequeued_url
wpull_hook.callbacks.handle_pre_response = handle_pre_response
wpull_hook.callbacks.handle_response = handle_response
wpull_hook.callbacks.handle_error = handle_error
wpull_hook.callbacks.get_urls = get_urls
wpull_hook.callbacks.wait_time = wait_time
# wpull_hook.callbacks.finish_statistics = finish_statistics
# wpull_hook.callbacks.exit_status = exit_status
wpull_hook.callbacks.version = 3
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
#  prestans, a standards based WSGI compliant REST framework for Python
#  http://prestans.org
#
#  Copyright (c) 2014, Eternity Technologies Pty Ltd.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of Eternity Technologies nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL ETERNITY TECHNOLOGIES BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

__all__ = [
    'UnsupportedVocabularyError',
    'UnsupportedContentTypeError',
    'ValidationError',
    'DataValidationException',
    'ParserException',
    'HandlerException',

    #: Parser Exceptions
    'NoEndpointError',
    'SerializationFailed',

    #: Data Validation
    'RequiredAttributeError',
    'ParseFailedError',
    'InvalidValue',
    'LessThanMinimum',
    'MoreThanMaximum',
    'InvalidChoice',
    'UnacceptableLength',
    'InvalidType',
    'InvalidCollection',
    'MissingParameter',
    'InvalidFormat',
    'InvalidMetaValue',
    'UnregisteredAdapter',

    #: Handler exceptions
    'ServiceUnavailable',
    'BadRequest',
    'Conflict',
    'NotFound',
    'Unauthorized',
    'MovedPermanently',
    'PaymentRequired',
    'Forbidden'
]

import prestans
import prestans.http

class Base(Exception):

    def __init__(self, http_status, message):

        self._http_status = http_status
        self._message = message
        self._stack_trace = list()

    @property
    def http_status(self):
        return self._http_status

    @http_status.setter
    def http_status(self, value):
        self._http_status = value

    @property
    def stack_trace(self):
        return self._stack_trace

    def push_trace(self, trace_object):
        self._stack_trace.append(trace_object)

    @property
    def message(self):
        return self._message

    def __unicode__(self):
        return unicode(self._message)

    def __str__(self):
        return unicode(self._message).encode('utf-8')

#:
#: These are top level exceptions that layout tell prestans how
#: the resulting messages are written out to the either the client
#: or the logger
#:
#: None of these exceptions are used directly, look at the sections
#: for usable exceptions.
#:
#: According to PEP 008 all exceptions that are error should have
#: the Error suffix e.g UnsupportedVocabularyError
#: http://www.python.org/dev/peps/pep-0008/#exception-names
#:

class UnsupportedVocabularyError(Base):
    """
    Called if none of the requested vocabularies are supported by the API
    this error message is sent as an HTML document. This exception is the
    only one that breaks serialisation rules.
    """

    def __init__(self, accept_header, supported_types):

        _code = prestans.http.STATUS.NOT_IMPLEMENTED
        _message = "Unsupported vocabulary in the Accept header"
        super(UnsupportedVocabularyError, self).__init__(_code, _message)

        self.push_trace({
            'accept_header': str(accept_header),
            'supported_types': supported_types
            })

class UnsupportedContentTypeError(Base):

    def __init__(self, requested_mime_type, content_type):

        _code = prestans.http.STATUS.NOT_IMPLEMENTED
        _message = "Unsupported Content-Type in Request"
        super(UnsupportedContentTypeError, self).__init__(_code, _message)

        self.push_trace({
            'requested_type': requested_mime_type,
            'supported_types': content_type
            })

class ValidationError(Base):
    """
    DataValidationException are Exceptions raised if prestans fails
    to validate data inbound or out using rules defined in Models.

    These contains a stack trace to indicate the depth of the error.
    E.g Model has Sub Model that has attribute which failed to validate.

    Each exception uses an HTTP status code and is sent to the client.
    """
    def __init__(self, message, attribute_name, value, blueprint):

        super(ValidationError, self).__init__(prestans.http.STATUS.BAD_REQUEST, message)
        self._attribute_name = attribute_name
        self._value = value

        self.append_validation_trace(blueprint)

    def append_validation_trace(self, blueprint):
        self.push_trace({
            'attribute_name': self._attribute_name,
            'value': self._value,
            'message': self._message,
            'blueprint': blueprint
            })

    #:
    #: The following are overridden to ensure sensible logging messages
    #:

    def __unicode__(self):
        _loggable_message = ("%s %s" % self._attribute_name)
        return unicode(self._loggable_message)

    def __str__(self):
        _loggable_message = ("%s %s" % self._attribute_name)
        return unicode(self._loggable_message).encode('utf-8')

class ParserException(Base):
    """
    ParserException are Exceptions raised if prestans fails to parse
    a request; these generally revolve around the Content-Types or
    missing payload. Specific parsing messages are of type DataValidationException
    """
    def __init__(self, code, message):
        super(ParserException, self).__init__(code, message)

class HandlerException(Base):
    """
    HandlerExceptions are Exceptions that are raised by handlers of the
    REST application. REST encourages the use different error codes to
    denote what the error is.

    E.g Use 404 to denote that a reqeusted doesn't exists on the server.

    If one of these exceptions do not match your user case; you are free
    construct your own error message and use an error code from outlined
    in prestans.http
    """
    def __init__(self, code, message):
        super(HandlerException, self).__init__(code, message)

#:
#: Parser Exception
#:

class UnimplementedVerbError(ParserException):

    def __init__(self, verb_name):

        _code = prestans.http.STATUS.NOT_IMPLEMENTED
        _message = "API does not implement the HTTP Verb"
        super(UnimplementedVerbError, self).__init__(_code, _message)

        self.push_trace({
            'verb': verb_name,
            })

class NoEndpointError(ParserException):

    def __init__(self):

        _code = prestans.http.STATUS.NOT_FOUND
        _message = "API does not provide this end-point"
        super(NoEndpointError, self).__init__(_code, _message)

class AuthenticationError(ParserException):

    def __init__(self, message=None):

        _code = prestans.http.STATUS.UNAUTHORIZED

        _message = message
        if _message is None:
            _message = "Authentication Error; service is only available to authenticated"

        super(AuthenticationError, self).__init__(_code, _message)

class AuthorizationError(ParserException):

    def __init__(self, role_name):

        _code = prestans.http.STATUS.FORBIDDEN
        _message = "%s is not allowed to access this resource"
        super(AuthorizationError, self).__init__(_code, _message)

class SerializationFailedError(ParserException):

    def __init__(self, format):

        _code = prestans.http.STATUS.NOT_FOUND
        _message = "Serialization failed"
        super(DeSerializationFailedError, self).__init__(_code, _message)

class DeSerializationFailedError(ParserException):

    def __init__(self, format):

        _code = prestans.http.STATUS.NOT_FOUND
        _message = "DeSerialization failed"
        super(DeSerializationFailedError, self).__init__(_code, _message)

class AttributeFilterDiffers(ParserException):
    """
    AttributeFilter initialised from request input does not conform to
    the configured template.
    """

    def __init__(self, attribute_list):

        _code = prestans.http.STATUS.BAD_REQUEST
        _message = "attribute filter has attributes that are not part of template"
        super(AttributeFilterDiffers, self).__init__(_code, _message)

        self.push_trace({
            'rejected_attribute_list': list(attribute_list),
            })

class InconsistentPersistentDataError(Base):
    """
    Raised by Data Adapters if validation of stored data fails. Extra information is
    written out to the server log and the client is returned a 500 response.

    Practical expectation is that this should not occur; however if REST rules have
    changed since the instantiation of the system.
    """

    def __init__(self, attribute_name, exception_message):
        _code = prestans.http.STATUS.INTERNAL_SERVER_ERROR
        _message = "Data Adapter failed to validate stored data on the server"
        super(InconsistentPersistentDataError, self).__init__(_code, _message)

        self._attribute_name = attribute_name

        self.push_trace({
            'exception_message': exception_message,
            'attribute_name': attribute_name
            })

    def __str__(self):
        return "DataAdapter failed to adapt %s, %s" % (self._attribute_name, self.message)

#:
#: Data Validation
#:

class DataValidationException(Base):

    def __init__(self, message):
        super(DataValidationException, self).__init__(prestans.http.STATUS.BAD_REQUEST, message)

class RequiredAttributeError(DataValidationException):

    def __init__(self):
        _message = "attribute is required and does not provide a default value"
        super(RequiredAttributeError, self).__init__(_message)

class ParseFailedError(DataValidationException):

    def __init__(self, message="Parser Failed"):
        _message = message
        super(ParseFailedError, self).__init__(_message)

class LessThanMinimumError(DataValidationException):

    def __init__(self, value, allowed_min):
        _message = "%i is less than the allowed minimum of %i" % (value, allowed_min)
        super(LessThanMinimumError, self).__init__(_message)

class MoreThanMaximumError(DataValidationException):

    def __init__(self, value, allowed_max):
        _message = "%i is more than the allowed maximum" % (value, allowed_max)
        super(MoreThanMaximumError, self).__init__(_message)

class InvalidChoiceError(DataValidationException):

    def __init__(self, value, allowed_choices):
        _message = "value is not one of these choics %s" % str(allowed_choices).strip('[]')
        super(InvalidChoiceError, self).__init__(_message)

class UnacceptableLengthError(DataValidationException):

    def __init__(self, value, minimum, maximum):
        _message = "value has to be %i and %i" % (minimum, maximum)
        super(UnacceptableLengthError, self).__init__(_message)

class InvalidType(DataValidationException):

    def __init__(self, value, type_name):
        _message = "invalid data type %s" % type_name
        super(InvalidType, self).__init__(_message)

class MissingParameterError(DataValidationException):

    def __init__(self):
        _message = "missing parameter"
        super(MissingParameterError, self).__init__(_message)

class InvalidFormatError(DataValidationException):

    def __init__(self, value):
        _message = "message goes in here"
        super(InvalidFormatError, self).__init__(_message)

class InvalidMetaValueError(DataValidationException):

    def __init__(self):
        _message = "message goes in here"
        super(InvalidMetaValueError, self).__init__(_message)

class UnregisteredAdapterError(DataValidationException):

    def __init__(self):
        _message = "message goes in here"
        super(UnregisteredAdapterError, self).__init__(_message)

#:
#: The following excepetions are used by REST handlers to indicate
#: commonly defined scenarios when dealing with data. BaseHandler traps
#: these and generates an appropriate error message for the end user.
#:

class ServiceUnavailable(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.SERVICE_UNAVAILABLE
        super(ServiceUnavailable, self).__init__(_code, message)

class BadRequest(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.BAD_REQUEST
        super(BadRequest, self).__init__(_code, message)

class Conflict(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.CONFLICT
        super(Conflict, self).__init__(_code, message)

class NotFound(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.NOT_FOUND
        super(NotFound, self).__init__(_code, message)

class Unauthorized(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.UNAUTHORIZED
        super(Unauthorized, self).__init__(_code, message)

class MovedPermanently(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.MOVED_PERMANENTLY
        super(MovedPermanently, self).__init__(_code, message)

class PaymentRequired(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.PAYMENT_REQUIRED
        super(PaymentRequired, self).__init__(_code, message)

class Forbidden(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.FORBIDDEN
        super(Forbidden, self).__init__(_code, message)
<commit_msg>Reference to the logging message<commit_after># -*- coding: utf-8 -*-
#
#  prestans, a standards based WSGI compliant REST framework for Python
#  http://prestans.org
#
#  Copyright (c) 2014, Eternity Technologies Pty Ltd.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of Eternity Technologies nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL ETERNITY TECHNOLOGIES BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

__all__ = [
    'UnsupportedVocabularyError',
    'UnsupportedContentTypeError',
    'ValidationError',
    'DataValidationException',
    'ParserException',
    'HandlerException',

    #: Parser Exceptions
    'NoEndpointError',
    'SerializationFailed',

    #: Data Validation
    'RequiredAttributeError',
    'ParseFailedError',
    'InvalidValue',
    'LessThanMinimum',
    'MoreThanMaximum',
    'InvalidChoice',
    'UnacceptableLength',
    'InvalidType',
    'InvalidCollection',
    'MissingParameter',
    'InvalidFormat',
    'InvalidMetaValue',
    'UnregisteredAdapter',

    #: Handler exceptions
    'ServiceUnavailable',
    'BadRequest',
    'Conflict',
    'NotFound',
    'Unauthorized',
    'MovedPermanently',
    'PaymentRequired',
    'Forbidden'
]

import prestans
import prestans.http

class Base(Exception):

    def __init__(self, http_status, message):

        self._http_status = http_status
        self._message = message
        self._stack_trace = list()

    @property
    def http_status(self):
        return self._http_status

    @http_status.setter
    def http_status(self, value):
        self._http_status = value

    @property
    def stack_trace(self):
        return self._stack_trace

    def push_trace(self, trace_object):
        self._stack_trace.append(trace_object)

    @property
    def message(self):
        return self._message

    def __unicode__(self):
        return unicode(self._message)

    def __str__(self):
        return unicode(self._message).encode('utf-8')

#:
#: These are top level exceptions that layout tell prestans how
#: the resulting messages are written out to the either the client
#: or the logger
#:
#: None of these exceptions are used directly, look at the sections
#: for usable exceptions.
#:
#: According to PEP 008 all exceptions that are error should have
#: the Error suffix e.g UnsupportedVocabularyError
#: http://www.python.org/dev/peps/pep-0008/#exception-names
#:

class UnsupportedVocabularyError(Base):
    """
    Called if none of the requested vocabularies are supported by the API
    this error message is sent as an HTML document. This exception is the
    only one that breaks serialisation rules.
    """

    def __init__(self, accept_header, supported_types):

        _code = prestans.http.STATUS.NOT_IMPLEMENTED
        _message = "Unsupported vocabulary in the Accept header"
        super(UnsupportedVocabularyError, self).__init__(_code, _message)

        self.push_trace({
            'accept_header': str(accept_header),
            'supported_types': supported_types
            })

class UnsupportedContentTypeError(Base):

    def __init__(self, requested_mime_type, content_type):

        _code = prestans.http.STATUS.NOT_IMPLEMENTED
        _message = "Unsupported Content-Type in Request"
        super(UnsupportedContentTypeError, self).__init__(_code, _message)

        self.push_trace({
            'requested_type': requested_mime_type,
            'supported_types': content_type
            })

class ValidationError(Base):
    """
    DataValidationException are Exceptions raised if prestans fails
    to validate data inbound or out using rules defined in Models.

    These contains a stack trace to indicate the depth of the error.
    E.g Model has Sub Model that has attribute which failed to validate.

    Each exception uses an HTTP status code and is sent to the client.
    """
    def __init__(self, message, attribute_name, value, blueprint):

        super(ValidationError, self).__init__(prestans.http.STATUS.BAD_REQUEST, message)
        self._attribute_name = attribute_name
        self._value = value

        self.append_validation_trace(blueprint)

    def append_validation_trace(self, blueprint):
        self.push_trace({
            'attribute_name': self._attribute_name,
            'value': self._value,
            'message': self._message,
            'blueprint': blueprint
            })

    #:
    #: The following are overridden to ensure sensible logging messages
    #:

    def __unicode__(self):
        _loggable_message = ("%s %s" % self._attribute_name)
        return unicode(_loggable_message)

    def __str__(self):
        _loggable_message = ("%s %s" % self._attribute_name)
        return unicode(_loggable_message).encode('utf-8')

class ParserException(Base):
    """
    ParserException are Exceptions raised if prestans fails to parse
    a request; these generally revolve around the Content-Types or
    missing payload. Specific parsing messages are of type DataValidationException
    """
    def __init__(self, code, message):
        super(ParserException, self).__init__(code, message)

class HandlerException(Base):
    """
    HandlerExceptions are Exceptions that are raised by handlers of the
    REST application. REST encourages the use different error codes to
    denote what the error is.

    E.g Use 404 to denote that a reqeusted doesn't exists on the server.

    If one of these exceptions do not match your user case; you are free
    construct your own error message and use an error code from outlined
    in prestans.http
    """
    def __init__(self, code, message):
        super(HandlerException, self).__init__(code, message)

#:
#: Parser Exception
#:

class UnimplementedVerbError(ParserException):

    def __init__(self, verb_name):

        _code = prestans.http.STATUS.NOT_IMPLEMENTED
        _message = "API does not implement the HTTP Verb"
        super(UnimplementedVerbError, self).__init__(_code, _message)

        self.push_trace({
            'verb': verb_name,
            })

class NoEndpointError(ParserException):

    def __init__(self):

        _code = prestans.http.STATUS.NOT_FOUND
        _message = "API does not provide this end-point"
        super(NoEndpointError, self).__init__(_code, _message)

class AuthenticationError(ParserException):

    def __init__(self, message=None):

        _code = prestans.http.STATUS.UNAUTHORIZED

        _message = message
        if _message is None:
            _message = "Authentication Error; service is only available to authenticated"

        super(AuthenticationError, self).__init__(_code, _message)

class AuthorizationError(ParserException):

    def __init__(self, role_name):

        _code = prestans.http.STATUS.FORBIDDEN
        _message = "%s is not allowed to access this resource"
        super(AuthorizationError, self).__init__(_code, _message)

class SerializationFailedError(ParserException):

    def __init__(self, format):

        _code = prestans.http.STATUS.NOT_FOUND
        _message = "Serialization failed"
        super(DeSerializationFailedError, self).__init__(_code, _message)

class DeSerializationFailedError(ParserException):

    def __init__(self, format):

        _code = prestans.http.STATUS.NOT_FOUND
        _message = "DeSerialization failed"
        super(DeSerializationFailedError, self).__init__(_code, _message)

class AttributeFilterDiffers(ParserException):
    """
    AttributeFilter initialised from request input does not conform to
    the configured template.
    """

    def __init__(self, attribute_list):

        _code = prestans.http.STATUS.BAD_REQUEST
        _message = "attribute filter has attributes that are not part of template"
        super(AttributeFilterDiffers, self).__init__(_code, _message)

        self.push_trace({
            'rejected_attribute_list': list(attribute_list),
            })

class InconsistentPersistentDataError(Base):
    """
    Raised by Data Adapters if validation of stored data fails. Extra information is
    written out to the server log and the client is returned a 500 response.

    Practical expectation is that this should not occur; however if REST rules have
    changed since the instantiation of the system.
    """

    def __init__(self, attribute_name, exception_message):
        _code = prestans.http.STATUS.INTERNAL_SERVER_ERROR
        _message = "Data Adapter failed to validate stored data on the server"
        super(InconsistentPersistentDataError, self).__init__(_code, _message)

        self._attribute_name = attribute_name

        self.push_trace({
            'exception_message': exception_message,
            'attribute_name': attribute_name
            })

    def __str__(self):
        return "DataAdapter failed to adapt %s, %s" % (self._attribute_name, self.message)

#:
#: Data Validation
#:

class DataValidationException(Base):

    def __init__(self, message):
        super(DataValidationException, self).__init__(prestans.http.STATUS.BAD_REQUEST, message)

class RequiredAttributeError(DataValidationException):

    def __init__(self):
        _message = "attribute is required and does not provide a default value"
        super(RequiredAttributeError, self).__init__(_message)

class ParseFailedError(DataValidationException):

    def __init__(self, message="Parser Failed"):
        _message = message
        super(ParseFailedError, self).__init__(_message)

class LessThanMinimumError(DataValidationException):

    def __init__(self, value, allowed_min):
        _message = "%i is less than the allowed minimum of %i" % (value, allowed_min)
        super(LessThanMinimumError, self).__init__(_message)

class MoreThanMaximumError(DataValidationException):

    def __init__(self, value, allowed_max):
        _message = "%i is more than the allowed maximum" % (value, allowed_max)
        super(MoreThanMaximumError, self).__init__(_message)

class InvalidChoiceError(DataValidationException):

    def __init__(self, value, allowed_choices):
        _message = "value is not one of these choics %s" % str(allowed_choices).strip('[]')
        super(InvalidChoiceError, self).__init__(_message)

class UnacceptableLengthError(DataValidationException):

    def __init__(self, value, minimum, maximum):
        _message = "value has to be %i and %i" % (minimum, maximum)
        super(UnacceptableLengthError, self).__init__(_message)

class InvalidType(DataValidationException):

    def __init__(self, value, type_name):
        _message = "invalid data type %s" % type_name
        super(InvalidType, self).__init__(_message)

class MissingParameterError(DataValidationException):

    def __init__(self):
        _message = "missing parameter"
        super(MissingParameterError, self).__init__(_message)

class InvalidFormatError(DataValidationException):

    def __init__(self, value):
        _message = "message goes in here"
        super(InvalidFormatError, self).__init__(_message)

class InvalidMetaValueError(DataValidationException):

    def __init__(self):
        _message = "message goes in here"
        super(InvalidMetaValueError, self).__init__(_message)

class UnregisteredAdapterError(DataValidationException):

    def __init__(self):
        _message = "message goes in here"
        super(UnregisteredAdapterError, self).__init__(_message)

#:
#: The following excepetions are used by REST handlers to indicate
#: commonly defined scenarios when dealing with data. BaseHandler traps
#: these and generates an appropriate error message for the end user.
#:

class ServiceUnavailable(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.SERVICE_UNAVAILABLE
        super(ServiceUnavailable, self).__init__(_code, message)

class BadRequest(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.BAD_REQUEST
        super(BadRequest, self).__init__(_code, message)

class Conflict(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.CONFLICT
        super(Conflict, self).__init__(_code, message)

class NotFound(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.NOT_FOUND
        super(NotFound, self).__init__(_code, message)

class Unauthorized(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.UNAUTHORIZED
        super(Unauthorized, self).__init__(_code, message)

class MovedPermanently(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.MOVED_PERMANENTLY
        super(MovedPermanently, self).__init__(_code, message)

class PaymentRequired(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.PAYMENT_REQUIRED
        super(PaymentRequired, self).__init__(_code, message)

class Forbidden(HandlerException):

    def __init__(self, message):
        _code = prestans.http.STATUS.FORBIDDEN
        super(Forbidden, self).__init__(_code, message)
<|endoftext|>


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

#
# Base functionality to convert XLS(x) files to CSV.
# Options to consider:
#   1. -f <filename>
#   2. -d <directory with files>
#       A. Implement a file search for
#

import sys
import csv
import threading
import mimetypes
import glob
import os

# If xlrd isn't installed, stop.
try:
    import xlrd
except ImportError, xlrd_error:
    print xlrd_error.message
    sys.exit(1)

def find_excel_files_by_mimetype():
    """
   Return a list of Excel files by matching MIMETYPE.
   :return: list
   """
    return [f for f in glob.glob("*.*")
        if mimetypes.guess_type(f)[0] == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet']

def find_excel_files_by_extension(extension, directory):
    """
   Return a list of Excel files by extension.
   :param extension: str
   :param directory: str
   :return: list
   """
    os.chdir(directory)
    return [f for f in glob.glob("*.{0}".format(extension))]

class WorkbookData(object):
    """
   Primary class to handle the provided workbook.
   """
    def __init__(self, spreadsheet):
        self._ss = xlrd.open_workbook(spreadsheet)
        self.sheets = self._ss.sheets()
        self.worksheet_count = self._ss.nsheets

        if self.worksheet_count == 1:
            self.worksheet_stats = {'columns': self._ss.sheets()[0].ncols,
                                    'rows': self._ss.sheets()[0].nrows}
        else:
            self.worksheet_stats = [{'name': sheet.name,
                                     'columns': sheet.ncols,
                                     'rows': sheet.nrows} for sheet in self.sheets]

class SaveToCSV(threading.Thread):
    """
   Save each sheet in the workbook to it's own CSV file.  Each
   worksheet is saved to CSV by it's own thread.
   """
    def __init__(self, workbook):
        threading.Thread.__init__(self)
        self.workbook = workbook

    def run(self):
        """
       Main worker.
       :return: None
       """
        for sheet in self.workbook.sheets:
            with open("{0}.csv".format(sheet.name), 'wb') as csv_file:
                cw = csv.writer(csv_file, quoting=csv.QUOTE_ALL)
                [cw.writerow(sheet.row_values(row_num))
                 for row_num in xrange(len(self.workbook.worksheet_stats))]

if __name__ == '__main__':
    wb = WorkbookData('test_file_1.xlsx')
    threads = []
    for num in xrange(wb.worksheet_count):
        thread = SaveToCSV(wb)
        thread.start()
        threads.append(thread)

    [t.join() for t in threads]

    if wb.worksheet_count == 1:
        print """Converted workbooks : {0}
             Rows      : {1}
             Colums    : {2}""".format(wb.worksheet_count,
                                        wb.worksheet_stats['rows'],
                                        wb.worksheet_stats['columns'])
    else:
        print "Converted workbooks : {0}".format(wb.worksheet_count)
        for sheet_data in wb.worksheet_stats:
            print "          Sheet     : {0}".format(sheet_data['name'])
            print """          Rows      : {0}
         Columns   : {2}\n""".format(wb.worksheet_count,
                                      sheet_data['rows'],
                                      sheet_data['columns'])
<commit_msg>minor changes<commit_after>#!/usr/bin/env python
# coding=utf-8

#
# Base functionality to convert XLS(x) files to CSV.
#

import os
import csv
import sys
import glob
import threading
import mimetypes

# If xlrd isn't installed, stop.
try:
    import xlrd
except ImportError, xlrd_error:
    print xlrd_error.message
    sys.exit(1)

def find_excel_files_by_mimetype():
    """
   Return a list of Excel files by matching MIMETYPE.
   :return: list
   """
    return [f for f in glob.glob("*.*")
        if mimetypes.guess_type(f)[0] == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet']

def find_excel_files_by_extension(extension, directory):
    """
   Return a list of Excel files by extension.
   :param extension: str
   :param directory: str
   :return: list
   """
    os.chdir(directory)
    return [f for f in glob.glob("*.{0}".format(extension))]

class WorkbookData(object):
    """
   Primary class to handle the provided workbook.
   """
    def __init__(self, spreadsheet):
        self._ss = xlrd.open_workbook(spreadsheet)
        self.sheets = self._ss.sheets()
        self.worksheet_count = self._ss.nsheets

        if self.worksheet_count == 1:
            self.worksheet_stats = {'columns': self._ss.sheets()[0].ncols,
                                    'rows': self._ss.sheets()[0].nrows}
        else:
            self.worksheet_stats = [{'name': sheet.name,
                                     'columns': sheet.ncols,
                                     'rows': sheet.nrows} for sheet in self.sheets]

class SaveToCSV(threading.Thread):
    """
   Save each sheet in the workbook to it's own CSV file.  Each
   worksheet is saved to CSV by it's own thread.
   """
    def __init__(self, workbook):
        threading.Thread.__init__(self)
        self.workbook = workbook

    def run(self):
        """
       Main worker.
       :return: None
       """
        for sheet in self.workbook.sheets:
            with open("{0}.csv".format(sheet.name), 'wb') as csv_file:
                cw = csv.writer(csv_file, quoting=csv.QUOTE_ALL)
                [cw.writerow(sheet.row_values(row_num))
                 for row_num in xrange(len(self.workbook.worksheet_stats))]

if __name__ == '__main__':
    wb = WorkbookData('test_file_1.xlsx')
    threads = []
    for num in xrange(wb.worksheet_count):
        thread = SaveToCSV(wb)
        thread.start()
        threads.append(thread)

    [t.join() for t in threads]

    if wb.worksheet_count == 1:
        print """Converted workbooks : {0}
             Rows      : {1}
             Colums    : {2}""".format(wb.worksheet_count,
                                       wb.worksheet_stats['rows'],
                                       wb.worksheet_stats['columns'])
    else:
        print "Converted workbooks : {0}".format(wb.worksheet_count)
        for sheet_data in wb.worksheet_stats:
            print "          Sheet     : {0}".format(sheet_data['name'])
            print """          Rows      : {0}
         Columns   : {2}\n""".format(wb.worksheet_count,
                                     sheet_data['rows'],
                                     sheet_data['columns'])
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2008 Tiny SPRL (<http://tiny.be>). All Rights Reserved
#    $Id$
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

pool_dic = {}

def get_db_and_pool(db_name, force_demo=False, status=None, update_module=False):
    if not status:
        status={}

    db = get_db_only(db_name)

    if db_name in pool_dic:
        pool = pool_dic[db_name]
    else:
        import addons
        import osv.osv
        pool = osv.osv.osv_pool()
        pool_dic[db_name] = pool
        addons.load_modules(db, force_demo, status, update_module)
        cr = db.cursor()
        pool.init_set(cr, False)
        cr.commit()
        cr.close()

        if not update_module:
            import report
            report.interface.register_all(db)
            pool.get('ir.cron')._poolJobs(db.dbname)
    return db, pool

def restart_pool(db_name, force_demo=False, update_module=False):
    del pool_dic[db_name]
    return get_db_and_pool(db_name, force_demo, update_module=update_module)

def get_db_only(db_name):
    # ATTENTION:
    # do not put this import outside this function
    # sql_db must not be loaded before the logger is initialized.
    # sql_db import psycopg2.tool which create a default logger if there is not.
    # this resulting of having the logs outputed twice...
    import sql_db
    db = sql_db.db_connect(db_name)
    return db

def get_db(db_name):
    return get_db_and_pool(db_name)[0]

def get_pool(db_name, force_demo=False, status=None, update_module=False):
    pool = get_db_and_pool(db_name, force_demo, status, update_module)[1]
    return pool

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:

<commit_msg>[FIX] Close the cursor if open<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2008 Tiny SPRL (<http://tiny.be>). All Rights Reserved
#    $Id$
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

pool_dic = {}

def get_db_and_pool(db_name, force_demo=False, status=None, update_module=False):
    if not status:
        status={}

    db = get_db_only(db_name)

    if db_name in pool_dic:
        pool = pool_dic[db_name]
    else:
        import addons
        import osv.osv
        pool = osv.osv.osv_pool()
        pool_dic[db_name] = pool
        addons.load_modules(db, force_demo, status, update_module)
        cr = db.cursor()
        try:
            pool.init_set(cr, False)
            cr.commit()
        finally:
            cr.close()

        if not update_module:
            import report
            report.interface.register_all(db)
            pool.get('ir.cron')._poolJobs(db.dbname)
    return db, pool

def restart_pool(db_name, force_demo=False, update_module=False):
    del pool_dic[db_name]
    return get_db_and_pool(db_name, force_demo, update_module=update_module)

def get_db_only(db_name):
    # ATTENTION:
    # do not put this import outside this function
    # sql_db must not be loaded before the logger is initialized.
    # sql_db import psycopg2.tool which create a default logger if there is not.
    # this resulting of having the logs outputed twice...
    import sql_db
    db = sql_db.db_connect(db_name)
    return db

def get_db(db_name):
    return get_db_and_pool(db_name)[0]

def get_pool(db_name, force_demo=False, status=None, update_module=False):
    pool = get_db_and_pool(db_name, force_demo, status, update_module)[1]
    return pool

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:

<|endoftext|>


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

import os
import re
import subprocess
import sys
import argparse
import glob

# core class
# {{{1

class Hub:
    DATA_FILE_TEMPLATE = '''\
# vim: noexpandtab list listchars=tab\:▸-

# NOTE:
# - Use TAB characters to separate titles and commands.
# - Comment lines (start with a '#') and empty lines are ignored

Edit fzf-cmdhub data file\t\t${EDITOR:-vi} ~/.fzf-cmdhub
'''

    MENU_PATH = os.path.expanduser('~/.fzf-cmdhub-menu')
    JOBS_DIR = os.path.expanduser('~/.fzf-cmdhub-jobs')
    AUTOLOAD_DIR = JOBS_DIR + '/autoload'

    # one or more Tabs separated line
    MENU_ITEM_PAT = r'^[^\t]+\t+[^\t]+$'
    SEP_PAT = r'\t+'

    def __init__(self):
        # check menu file avalability
        # if not exists, create & populate it with initial content
        if not os.path.exists(self.MENU_PATH):
            with open(self.MENU_PATH, 'w') as f:
                f.write(self.DATA_FILE_TEMPLATE)

        if not os.path.exists(self.JOBS_DIR):
            subprocess.call('mkdir -p' + self.AUTOLOAD_DIR, shell=True)

        with open(self.MENU_PATH, 'r') as data_file:
            lines = [
                l for l in data_file if re.match(
                    self.MENU_ITEM_PAT, l)]

            title_cmd_pairs = map(
                lambda l: re.split(
                    self.SEP_PAT, l), lines)

            title_cmd_pairs += self.autoload_pairs()

            # check for duplicate title
            sorted_titles = sorted([p[0] for p in title_cmd_pairs])
            for i in range(len(sorted_titles) - 1):
                if sorted_titles[i] == sorted_titles[i + 1]:
                    sys.stderr.write(
                        '* Found duplicate title: {} *\n'.format(sorted_titles[i]))

            # handle sharp lines
            def translate_sharp_line(line):
                line = line.strip()

                if line.startswith('#s '):
                    line = re.sub(
                        r'^#s\s+', 'source {}/'.format(self.JOBS_DIR),
                        line)
                elif line.startswith('#e '):
                    line = re.sub(
                        r'^#e\s+', 'exec {}/'.format(self.JOBS_DIR),
                        line)
                elif line.startswith('#x '):
                    line = re.sub(
                        r'^#x\s+', '{}/'.format(self.JOBS_DIR),
                        line)

                return line

            title_cmd_pairs = map(
                lambda p: (p[0], translate_sharp_line(p[1])),
                title_cmd_pairs)

            self.core_dict = dict(title_cmd_pairs)

    def autoload_pairs(self):
        """
        parse files under .fzf-cmdhub-jobs/autoload
        return: a list of 2-tuple (title, cmd_line)
        """

        title_cmd_pairs = []

        files = glob.glob(self.AUTOLOAD_DIR + '/*.[bex]')
        if len(files) == 0:
            return []

        for fn in files:
            # read first 2 lines
            # title line must stay in the first 2 lines
            with open(fn) as fp:
                first_2_lines = [fp.readline(), fp.readline()]

            # collect title if any
            # sharp must start with '#cmdhub:' or '# cmdhub:'
            for l in first_2_lines:
                if re.match(u'#\s*cmdhub:', l):
                    title = l[l.index(':') + 1:].strip()
                    break
            else:
                os.stderr.write('* can not find cmdhub title line in first 2' +
                                'lines of autoload/{} *'.format(fn))
                break

            pair = (title,
                    '#{} autoload/{}'.format(fn[-1],
                                             os.path.basename(fn)))
            title_cmd_pairs.append(pair)

        return title_cmd_pairs

    def print_titles(self):
        print('\n'.join(self.core_dict.keys()))

    def print_cmd_for_title(self, title):
        print(self.core_dict[title])
# }}}1

# command interface
# {{{1
ap = argparse.ArgumentParser(
    prog='fzf-cmdhub',
    description='A fzf extension: cmdhub',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter
)

grp = ap.add_mutually_exclusive_group(required=True)

grp.add_argument(
    '-t',
    '--list-titles',
    action='store_true',
    dest='action',
)

grp.add_argument(
    '-c',
    '--cmd-for',
    metavar='TITLE',
    dest='action',
)

ns = ap.parse_args()

the_hub = Hub()
if ns.action is True:
    the_hub.print_titles()
else:
    the_hub.print_cmd_for_title(ns.action)

# }}}1

# for test
# the_hub = Hub()
# print('\n'.join(the_hub.core_dict.values()))

# vim: foldmethod=marker
<commit_msg>fzf-cmd_line.py: fix bug, [bex] -> [sex]<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import re
import subprocess
import sys
import argparse
import glob

# CORE CLASS                                                                           {{{1

class Hub:
    DATA_FILE_TEMPLATE = '''\
# vim: noexpandtab list listchars=tab\:▸-

# NOTE:
# - Use TAB characters to separate titles and commands.
# - Comment lines (start with a '#') and empty lines are ignored

Edit fzf-cmdhub data file\t\t${EDITOR:-vi} ~/.fzf-cmdhub
'''

    MENU_PATH = os.path.expanduser('~/.fzf-cmdhub-menu')
    JOBS_DIR = os.path.expanduser('~/.fzf-cmdhub-jobs')
    AUTOLOAD_DIR = JOBS_DIR + '/autoload'

    # one or more Tabs separated line
    MENU_ITEM_PAT = r'^[^\t]+\t+[^\t]+$'
    SEP_PAT = r'\t+'

    def __init__(self):
        # check menu file avalability
        # if not exists, create & populate it with initial content
        if not os.path.exists(self.MENU_PATH):
            with open(self.MENU_PATH, 'w') as f:
                f.write(self.DATA_FILE_TEMPLATE)

        if not os.path.exists(self.JOBS_DIR):
            subprocess.call('mkdir -p' + self.AUTOLOAD_DIR, shell=True)

        with open(self.MENU_PATH, 'r') as data_file:
            lines = [
                l for l in data_file if re.match(
                    self.MENU_ITEM_PAT, l)]

            title_cmd_pairs = map(
                lambda l: re.split(
                    self.SEP_PAT, l), lines)

            title_cmd_pairs += self.autoload_pairs()

            # check for duplicate title
            sorted_titles = sorted([p[0] for p in title_cmd_pairs])
            for i in range(len(sorted_titles) - 1):
                if sorted_titles[i] == sorted_titles[i + 1]:
                    sys.stderr.write(
                        '* Found duplicate title: {} *\n'.format(sorted_titles[i]))

            # handle sharp lines
            def translate_sharp_line(line):
                line = line.strip()

                if line.startswith('#s '):
                    line = re.sub(
                        r'^#s\s+', 'source {}/'.format(self.JOBS_DIR),
                        line)
                elif line.startswith('#e '):
                    line = re.sub(
                        r'^#e\s+', 'exec {}/'.format(self.JOBS_DIR),
                        line)
                elif line.startswith('#x '):
                    line = re.sub(
                        r'^#x\s+', '{}/'.format(self.JOBS_DIR),
                        line)

                return line

            title_cmd_pairs = map(
                lambda p: (p[0], translate_sharp_line(p[1])),
                title_cmd_pairs)

            self.core_dict = dict(title_cmd_pairs)

    def autoload_pairs(self):
        """
        parse files under .fzf-cmdhub-jobs/autoload
        return: a list of 2-tuple (title, cmd_line)
        """

        title_cmd_pairs = []

        files = glob.glob(self.AUTOLOAD_DIR + '/*.[sex]')
        if len(files) == 0:
            return []

        for fn in files:
            # read first 2 lines
            # title line must stay in the first 2 lines
            with open(fn) as fp:
                first_2_lines = [fp.readline(), fp.readline()]

            # collect title if any
            # sharp must start with '#cmdhub:' or '# cmdhub:'
            for l in first_2_lines:
                if re.match(u'#\s*cmdhub:', l):
                    title = l[l.index(':') + 1:].strip()
                    break
            else:
                os.stderr.write('* can not find cmdhub title line in first 2' +
                                'lines of autoload/{} *'.format(fn))
                break

            pair = (title,
                    '#{} autoload/{}'.format(fn[-1],
                                             os.path.basename(fn)))
            title_cmd_pairs.append(pair)

        return title_cmd_pairs

    def print_titles(self):
        print('\n'.join(self.core_dict.keys()))

    def print_cmd_for_title(self, title):
        print(self.core_dict[title])
# }}}1

# COMMAND INTERFACE                                                                    {{{1

ap = argparse.ArgumentParser(
    prog='fzf-cmdhub',
    description='A fzf extension: cmdhub',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter
)

grp = ap.add_mutually_exclusive_group(required=True)

grp.add_argument(
    '-t',
    '--list-titles',
    action='store_true',
    dest='action',
)

grp.add_argument(
    '-c',
    '--cmd-for',
    metavar='TITLE',
    dest='action',
)

ns = ap.parse_args()

the_hub = Hub()
if ns.action is True:
    the_hub.print_titles()
else:
    the_hub.print_cmd_for_title(ns.action)

# }}}1

# for test
# the_hub = Hub()
# print('\n'.join(the_hub.core_dict.values()))

# vim: foldmethod=marker
<|endoftext|>


<commit_before>import os.path
import re

from vial import vim, vfunc
from vial.utils import get_var, focus_window
from vial.widgets import ListFormatter, ListView, SearchDialog

from .. import quick_open as module
from . import search

module.dialog = None

IGNORE_DIRS = re.compile(r'(^|.*/)(\.git|\.svn|\.hg)$')
IGNORE_FILES = re.compile(r'^.*(\.pyc|\.pyo|\.swp|\.class|\.o)$')

def quick_open():
    if not module.dialog:
        module.dialog = QuickOpen()

    module.dialog.open()

class QuickOpen(SearchDialog):
    def __init__(self):
        self.filelist = []
        SearchDialog.__init__(self, '__vial_quick_open',
            ListView(self.filelist, ListFormatter(0, 0, 1, 1)))

    def open(self):
        self.last_window = vfunc.winnr()
        self.roots = get_var('vial_quick_open_projects', [os.getcwd()])
        self.list_view.clear()
        self.show(u'')

    def on_select(self, item, cursor):
        focus_window(self.last_window)
        vim.command('e {}'.format(item[2]))

    def on_prompt_changed(self, prompt):
        if prompt:
            self.loop.idle(self.fill(prompt))

    def fill(self, prompt):
        current = self.current = object()
        self.list_view.clear()
        last_index = 0
        cnt = 0
        already_matched = {}

        file_cache = {}
        def fill_cache(seq, cache):
            for r in seq:
                cache.append(r)
                yield r

        for m in search.get_matchers(prompt):
            for r in self.roots:
                if r in file_cache:
                    flist = file_cache[r]
                else:
                    cache = file_cache[r] = []
                    flist = fill_cache(search.get_files(r, '', IGNORE_FILES, IGNORE_DIRS), cache)

                for name, path, root, top, fpath in flist:
                    if current is not self.current:
                        return

                    if fpath not in already_matched and m(name, path):
                        already_matched[fpath] = True

                        self.filelist.append((name, top, fpath, root))
                        if len(self.filelist) > 20:
                            self.list_view.render()
                            return

                    cnt += 1
                    if not cnt % 50:
                        self.list_view.render()
                        self.loop.refresh()
                        yield

            self.list_view.render()

<commit_msg>explicit loop enter. on_cancel handler<commit_after>import os.path
import re

from vial import vim, vfunc
from vial.utils import get_var, focus_window
from vial.widgets import ListFormatter, ListView, SearchDialog

from .. import quick_open as module
from . import search

module.dialog = None

IGNORE_DIRS = re.compile(r'(^|.*/)(\.git|\.svn|\.hg)$')
IGNORE_FILES = re.compile(r'^.*(\.pyc|\.pyo|\.swp|\.class|\.o)$')

def quick_open():
    if not module.dialog:
        module.dialog = QuickOpen()

    module.dialog.open()

class QuickOpen(SearchDialog):
    def __init__(self):
        self.filelist = []
        SearchDialog.__init__(self, '__vial_quick_open',
            ListView(self.filelist, ListFormatter(0, 0, 1, 1)))

    def open(self):
        self.last_window = vfunc.winnr()
        self.roots = get_var('vial_quick_open_projects', [os.getcwd()])
        self.list_view.clear()
        self.show(u'')
        self.loop.enter()

    def on_select(self, item, cursor):
        focus_window(self.last_window)
        vim.command('e {}'.format(item[2]))

    def on_cancel(self):
        focus_window(self.last_window)

    def on_prompt_changed(self, prompt):
        if prompt:
            self.loop.idle(self.fill(prompt))
        else:
            self.list_view.show_cursor(False)
            self.buf[0:] = ['Type something to search']
            self.loop.refresh()

    def fill(self, prompt):
        current = self.current = object()
        self.list_view.clear()
        last_index = 0
        cnt = 0
        already_matched = {}

        file_cache = {}
        def fill_cache(seq, cache):
            for r in seq:
                cache.append(r)
                yield r

        for m in search.get_matchers(prompt):
            for r in self.roots:
                if r in file_cache:
                    flist = file_cache[r]
                else:
                    cache = file_cache[r] = []
                    flist = fill_cache(search.get_files(r, '', IGNORE_FILES, IGNORE_DIRS), cache)

                for name, path, root, top, fpath in flist:
                    if current is not self.current:
                        return

                    if fpath not in already_matched and m(name, path):
                        already_matched[fpath] = True

                        self.filelist.append((name, top, fpath, root))
                        if len(self.filelist) > 20:
                            self.list_view.render()
                            return

                    cnt += 1
                    if not cnt % 50:
                        self.list_view.render()
                        self.loop.refresh()
                        yield

            self.list_view.render()

<|endoftext|>