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

'''
DigitalOcean external inventory script
======================================

Generates Ansible inventory of DigitalOcean Droplets.

In addition to the --list and --host options used by Ansible, there are options
for generating JSON of other DigitalOcean data.  This is useful when creating
droplets.  For example, --regions will return all the DigitalOcean Regions.
This information can also be easily found in the cache file, whose default
location is /tmp/ansible-digital_ocean.cache).

The --pretty (-p) option pretty-prints the output for better human readability.

----
Although the cache stores all the information received from DigitalOcean,
the cache is not used for current droplet information (in --list, --host,
--all, and --droplets).  This is so that accurate droplet information is always
found.  You can force this script to use the cache with --force-cache.

----
Configuration is read from `digital_ocean.ini`, then from environment variables,
then and command-line arguments.

Most notably, the DigitalOcean Client ID and API Key must be specified.  They
can be specified in the INI file or with the following environment variables:
    export DO_CLIENT_ID='DO123' DO_API_KEY='abc123'

Alternatively, they can be passed on the command-line with --client-id and
--api-key.

If you specify DigitalOcean credentials in the INI file, a handy way to
get them into your environment (e.g., to use the digital_ocean module)
is to use the output of the --env option with export:
    export $(digital_ocean.py --env)

----
The following groups are generated from --list:
 - ID    (droplet ID)
 - NAME  (droplet NAME)
 - image_ID
 - image_NAME
 - distro_NAME  (distribution NAME from image)
 - region_ID
 - region_NAME
 - size_ID
 - size_NAME
 - status_STATUS

When run against a specific host, this script returns the following variables:
 - do_created_at
 - do_distroy
 - do_id
 - do_image
 - do_image_id
 - do_ip_address
 - do_name
 - do_region
 - do_region_id
 - do_size
 - do_size_id
 - do_status

-----
```
usage: digital_ocean.py [-h] [--list] [--host HOST] [--all]
                                 [--droplets] [--regions] [--images] [--sizes]
                                 [--ssh-keys] [--domains] [--pretty]
                                 [--cache-path CACHE_PATH]
                                 [--cache-max_age CACHE_MAX_AGE]
                                 [--refresh-cache] [--client-id CLIENT_ID]
                                 [--api-key API_KEY]

Produce an Ansible Inventory file based on DigitalOcean credentials

optional arguments:
  -h, --help            show this help message and exit
  --list                List all active Droplets as Ansible inventory
                        (default: True)
  --host HOST           Get all Ansible inventory variables about a specific
                        Droplet
  --all                 List all DigitalOcean information as JSON
  --droplets            List Droplets as JSON
  --regions             List Regions as JSON
  --images              List Images as JSON
  --sizes               List Sizes as JSON
  --ssh-keys            List SSH keys as JSON
  --domains             List Domains as JSON
  --pretty, -p          Pretty-print results
  --cache-path CACHE_PATH
                        Path to the cache files (default: .)
  --cache-max_age CACHE_MAX_AGE
                        Maximum age of the cached items (default: 0)
  --refresh-cache       Force refresh of cache by making API requests to
                        DigitalOcean (default: False - use cache files)
  --client-id CLIENT_ID, -c CLIENT_ID
                        DigitalOcean Client ID
  --api-key API_KEY, -a API_KEY
                        DigitalOcean API Key
```

'''

# (c) 2013, Evan Wies <evan@neomantra.net>
#
# Inspired by the EC2 inventory plugin:
# https://github.com/ansible/ansible/blob/devel/plugins/inventory/ec2.py
#
# This file is part of Ansible,
#
# Ansible 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.
#
# Ansible 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 Ansible.  If not, see <http://www.gnu.org/licenses/>.

######################################################################

import os
import sys
import re
import argparse
from time import time
import ConfigParser

try:
    import json
except ImportError:
    import simplejson as json

try:
    from dopy.manager import DoError, DoManager
except ImportError, e:
    print "failed=True msg='`dopy` library required for this script'"
    sys.exit(1)

class DigitalOceanInventory(object):

    ###########################################################################
    # Main execution path
    ###########################################################################

    def __init__(self):
        ''' Main execution path '''

        # DigitalOceanInventory data
        self.data = {}      # All DigitalOcean data
        self.inventory = {} # Ansible Inventory
        self.index = {}     # Various indices of Droplet metadata

        # Define defaults
        self.cache_path = '.'
        self.cache_max_age = 0

        # Read settings, environment variables, and CLI arguments
        self.read_settings()
        self.read_environment()
        self.read_cli_args()

        # Verify credentials were set
        if not hasattr(self, 'client_id') or not hasattr(self, 'api_key'):
            print '''Could not find values for DigitalOcean client_id and api_key.
They must be specified via either ini file, command line argument (--client-id and --api-key),
or environment variables (DO_CLIENT_ID and DO_API_KEY)'''
            sys.exit(-1)

        # env command, show DigitalOcean credentials
        if self.args.env:
            print "DO_CLIENT_ID=%s DO_API_KEY=%s" % (self.client_id, self.api_key)
            sys.exit(0)

        # Manage cache
        self.cache_filename = self.cache_path + "/ansible-digital_ocean.cache"
        self.cache_refreshed = False

        if not self.args.force_cache and self.args.refresh_cache or not self.is_cache_valid():
            self.load_all_data_from_digital_ocean()
        else:
            self.load_from_cache()
            if len(self.data) == 0:
                if self.args.force_cache:
                    print '''Cache is empty and --force-cache was specified'''
                    sys.exit(-1)
                self.load_all_data_from_digital_ocean()
            else:
                # We always get fresh droplets for --list, --host, --all, and --droplets
                # unless --force-cache is specified
                if not self.args.force_cache and (
                   self.args.list or self.args.host or self.args.all or self.args.droplets):
                    self.load_droplets_from_digital_ocean()

        # Pick the json_data to print based on the CLI command
        if self.args.droplets:   json_data = { 'droplets': self.data['droplets'] }
        elif self.args.regions:  json_data = { 'regions':  self.data['regions'] }
        elif self.args.images:   json_data = { 'images':   self.data['images'] }
        elif self.args.sizes:    json_data = { 'sizes':    self.data['sizes'] }
        elif self.args.ssh_keys: json_data = { 'ssh_keys': self.data['ssh_keys'] }
        elif self.args.domains:  json_data = { 'domains':  self.data['domains'] }
        elif self.args.all:      json_data = self.data

        elif self.args.host:     json_data = self.load_droplet_variables_for_host()
        else:    # '--list' this is last to make it default
                                 json_data = self.inventory

        if self.args.pretty:
            print json.dumps(json_data, sort_keys=True, indent=2)
        else:
            print json.dumps(json_data)
        # That's all she wrote...

    ###########################################################################
    # Script configuration
    ###########################################################################

    def read_settings(self):
        ''' Reads the settings from the digital_ocean.ini file '''
        config = ConfigParser.SafeConfigParser()
        config.read(os.path.dirname(os.path.realpath(__file__)) + '/digital_ocean.ini')

        # Credentials
        if config.has_option('digital_ocean', 'client_id'):
            self.client_id = config.get('digital_ocean', 'client_id')
        if config.has_option('digital_ocean', 'api_key'):
            self.api_key = config.get('digital_ocean', 'api_key')

        # Cache related
        if config.has_option('digital_ocean', 'cache_path'):
            self.cache_path = config.get('digital_ocean', 'cache_path')
        if config.has_option('digital_ocean', 'cache_max_age'):
            self.cache_max_age = config.getint('digital_ocean', 'cache_max_age')

    def read_environment(self):
        ''' Reads the settings from environment variables '''
        # Setup credentials
        if os.getenv("DO_CLIENT_ID"): self.client_id = os.getenv("DO_CLIENT_ID")
        if os.getenv("DO_API_KEY"):   self.api_key = os.getenv("DO_API_KEY")

    def read_cli_args(self):
        ''' Command line argument processing '''
        parser = argparse.ArgumentParser(description='Produce an Ansible Inventory file based on DigitalOcean credentials')

        parser.add_argument('--list', action='store_true', help='List all active Droplets as Ansible inventory (default: True)')
        parser.add_argument('--host', action='store', help='Get all Ansible inventory variables about a specific Droplet')

        parser.add_argument('--all', action='store_true', help='List all DigitalOcean information as JSON')
        parser.add_argument('--droplets','-d', action='store_true', help='List Droplets as JSON')
        parser.add_argument('--regions', action='store_true', help='List Regions as JSON')
        parser.add_argument('--images', action='store_true', help='List Images as JSON')
        parser.add_argument('--sizes', action='store_true', help='List Sizes as JSON')
        parser.add_argument('--ssh-keys', action='store_true', help='List SSH keys as JSON')
        parser.add_argument('--domains', action='store_true',help='List Domains as JSON')

        parser.add_argument('--pretty','-p', action='store_true', help='Pretty-print results')

        parser.add_argument('--cache-path', action='store', help='Path to the cache files (default: .)')
        parser.add_argument('--cache-max_age', action='store', help='Maximum age of the cached items (default: 0)')
        parser.add_argument('--force-cache', action='store_true', default=False, help='Only use data from the cache')
        parser.add_argument('--refresh-cache','-r', action='store_true', default=False, help='Force refresh of cache by making API requests to DigitalOcean (default: False - use cache files)')

        parser.add_argument('--env','-e', action='store_true', help='Display DO_CLIENT_ID and DO_API_KEY')
        parser.add_argument('--client-id','-c', action='store', help='DigitalOcean Client ID')
        parser.add_argument('--api-key','-a', action='store', help='DigitalOcean API Key')

        self.args = parser.parse_args()

        if self.args.client_id: self.client_id = self.args.client_id
        if self.args.api_key: self.api_key = self.args.api_key
        if self.args.cache_path: self.cache_path = self.args.cache_path
        if self.args.cache_max_age: self.cache_max_age = self.args.cache_max_age

        # Make --list default if none of the other commands are specified
        if (not self.args.droplets and not self.args.regions and not self.args.images and
            not self.args.sizes and not self.args.ssh_keys and not self.args.domains and
            not self.args.all and not self.args.host):
                self.args.list = True

    ###########################################################################
    # Data Management
    ###########################################################################

    def load_all_data_from_digital_ocean(self):
        ''' Use dopy to get all the information from DigitalOcean and save data in cache files '''
        manager  = DoManager(self.client_id, self.api_key)

        self.data = {}
        self.data['droplets'] = self.sanitize_list(manager.all_active_droplets())
        self.data['regions']  = self.sanitize_list(manager.all_regions())
        self.data['images']   = self.sanitize_list(manager.all_images(filter=None))
        self.data['sizes']    = self.sanitize_list(manager.sizes())
        self.data['ssh_keys'] = self.sanitize_list(manager.all_ssh_keys())
        self.data['domains']  = self.sanitize_list(manager.all_domains())

        self.index = {}
        self.index['region_to_name']  = self.build_index(self.data['regions'], 'id', 'name')
        self.index['size_to_name']    = self.build_index(self.data['sizes'], 'id', 'name')
        self.index['image_to_name']   = self.build_index(self.data['images'], 'id', 'name')
        self.index['image_to_distro'] = self.build_index(self.data['images'], 'id', 'distribution')
        self.index['host_to_droplet'] = self.build_index(self.data['droplets'], 'ip_address', 'id', False)

        self.build_inventory()

        self.write_to_cache()

    def load_droplets_from_digital_ocean(self):
        ''' Use dopy to get droplet information from DigitalOcean and save data in cache files '''
        manager  = DoManager(self.client_id, self.api_key)
        self.data['droplets'] = self.sanitize_list(manager.all_active_droplets())
        self.index['host_to_droplet'] = self.build_index(self.data['droplets'], 'ip_address', 'id', False)
        self.build_inventory()
        self.write_to_cache()

    def build_index(self, source_seq, key_from, key_to, use_slug=True):
        dest_dict = {}
        for item in source_seq:
            name = (use_slug and item.has_key('slug')) and item['slug'] or item[key_to]
            key = item[key_from]
            dest_dict[key] = name
        return dest_dict

    def build_inventory(self):
        '''Build Ansible inventory of droplets'''
        self.inventory = {}

        # add all droplets by id and name
        for droplet in self.data['droplets']:
            dest = droplet['ip_address']

            self.inventory[droplet['id']] = [dest]
            self.push(self.inventory, droplet['name'], dest)
            self.push(self.inventory, 'region_'+droplet['region_id'], dest)
            self.push(self.inventory, 'image_' +droplet['image_id'], dest)
            self.push(self.inventory, 'size_'  +droplet['size_id'], dest)
            self.push(self.inventory, 'status_'+droplet['status'], dest)

            region_name = self.index['region_to_name'].get(droplet['region_id'])
            if region_name:
                self.push(self.inventory, 'region_'+region_name, dest)

            size_name = self.index['size_to_name'].get(droplet['size_id'])
            if size_name:
                self.push(self.inventory, 'size_'+size_name, dest)

            image_name = self.index['image_to_name'].get(droplet['image_id'])
            if image_name:
                self.push(self.inventory, 'image_'+image_name, dest)

            distro_name = self.index['image_to_distro'].get(droplet['image_id'])
            if distro_name:
                self.push(self.inventory, 'distro_'+distro_name, dest)

    def load_droplet_variables_for_host(self):
        '''Generate a JSON response to a --host call'''
        host = self.to_safe(str(self.args.host))

        if not host in self.index['host_to_droplet']:
            # try updating cache
            if not self.args.force_cache:
                self.load_all_data_from_digital_ocean()
            if not host in self.index['host_to_droplet']:
                # host might not exist anymore
                return {}

        droplet = None
        if self.cache_refreshed:
            for drop in self.data['droplets']:
                if drop['ip_address'] == host:
                    droplet = self.sanitize_dict(drop)
                    break
        else:
            # Cache wasn't refreshed this run, so hit DigitalOcean API
            manager = DoManager(self.client_id, self.api_key)
            droplet_id = self.index['host_to_droplet'][host]
            droplet = self.sanitize_dict(manager.show_droplet(droplet_id))

        if not droplet:
            return {}

        # Put all the information in a 'do_' namespace
        info = {}
        for k, v in droplet.items():
            info['do_'+k] = v

        # Generate user-friendly variables (i.e. not the ID's)
        if droplet.has_key('region_id'):
            info['do_region'] = self.index['region_to_name'].get(droplet['region_id'])
        if droplet.has_key('size_id'):
            info['do_size'] = self.index['size_to_name'].get(droplet['size_id'])
        if droplet.has_key('image_id'):
            info['do_image']  = self.index['image_to_name'].get(droplet['image_id'])
            info['do_distro'] = self.index['image_to_distro'].get(droplet['image_id'])

        return info

    ###########################################################################
    # Cache Management
    ###########################################################################

    def is_cache_valid(self):
        ''' Determines if the cache files have expired, or if it is still valid '''
        if os.path.isfile(self.cache_filename):
            mod_time = os.path.getmtime(self.cache_filename)
            current_time = time()
            if (mod_time + self.cache_max_age) > current_time:
                return True
        return False

    def load_from_cache(self):
        ''' Reads the data from the cache file and assigns it to member variables as Python Objects'''
        cache = open(self.cache_filename, 'r')
        json_data = cache.read()
        cache.close()
        data = json.loads(json_data)

        self.data = data['data']
        self.inventory = data['inventory']
        self.index = data['index']

    def write_to_cache(self):
        ''' Writes data in JSON format to a file '''
        data = { 'data': self.data, 'index': self.index, 'inventory': self.inventory }
        json_data = json.dumps(data, sort_keys=True, indent=2)

        cache = open(self.cache_filename, 'w')
        cache.write(json_data)
        cache.close()

    ###########################################################################
    # Utilities
    ###########################################################################

    def push(self, my_dict, key, element):
        ''' Pushed an element onto an array that may not have been defined in the dict '''
        if key in my_dict:
            my_dict[key].append(element);
        else:
            my_dict[key] = [element]

    def to_safe(self, word):
        ''' Converts 'bad' characters in a string to underscores so they can be used as Ansible groups '''
        return re.sub("[^A-Za-z0-9\-\.]", "_", word)

    def sanitize_dict(self, d):
        new_dict = {}
        for k, v in d.items():
            if v != None:
                new_dict[self.to_safe(str(k))] = self.to_safe(str(v))
        return new_dict

    def sanitize_list(self, seq):
        new_seq = []
        for d in seq:
            new_seq.append(self.sanitize_dict(d))
        return new_seq

###########################################################################
# Run the script
DigitalOceanInventory()
<commit_msg>Update DigitalOcean dynamic inventory to API v2<commit_after>#!/usr/bin/env python

'''
DigitalOcean external inventory script
======================================

Generates Ansible inventory of DigitalOcean Droplets.

In addition to the --list and --host options used by Ansible, there are options
for generating JSON of other DigitalOcean data.  This is useful when creating
droplets.  For example, --regions will return all the DigitalOcean Regions.
This information can also be easily found in the cache file, whose default
location is /tmp/ansible-digital_ocean.cache).

The --pretty (-p) option pretty-prints the output for better human readability.

----
Although the cache stores all the information received from DigitalOcean,
the cache is not used for current droplet information (in --list, --host,
--all, and --droplets).  This is so that accurate droplet information is always
found.  You can force this script to use the cache with --force-cache.

----
Configuration is read from `digital_ocean.ini`, then from environment variables,
then and command-line arguments.

Most notably, the DigitalOcean Client ID and API Key must be specified.  They
can be specified in the INI file or with the following environment variables:
    export DO_CLIENT_ID='DO123' DO_API_KEY='abc123'

Alternatively, they can be passed on the command-line with --client-id and
--api-key.

If you specify DigitalOcean credentials in the INI file, a handy way to
get them into your environment (e.g., to use the digital_ocean module)
is to use the output of the --env option with export:
    export $(digital_ocean.py --env)

----
The following groups are generated from --list:
 - ID    (droplet ID)
 - NAME  (droplet NAME)
 - image_ID
 - image_NAME
 - distro_NAME  (distribution NAME from image)
 - region_ID
 - region_NAME
 - size_ID
 - size_NAME
 - status_STATUS

When run against a specific host, this script returns the following variables:
 - do_created_at
 - do_distroy
 - do_id
 - do_image
 - do_image_id
 - do_ip_address
 - do_name
 - do_region
 - do_region_id
 - do_size
 - do_size_id
 - do_status

-----
```
usage: digital_ocean.py [-h] [--list] [--host HOST] [--all]
                                 [--droplets] [--regions] [--images] [--sizes]
                                 [--ssh-keys] [--domains] [--pretty]
                                 [--api-token API_TOKEN]

Produce an Ansible Inventory file based on DigitalOcean credentials

optional arguments:
  -h, --help            show this help message and exit
  --list                List all active Droplets as Ansible inventory
                        (default: True)
  --host HOST           Get all Ansible inventory variables about a specific
                        Droplet
  --all                 List all DigitalOcean information as JSON
  --droplets            List Droplets as JSON
  --regions             List Regions as JSON
  --images              List Images as JSON
  --sizes               List Sizes as JSON
  --ssh-keys            List SSH keys as JSON
  --domains             List Domains as JSON
  --pretty, -p          Pretty-print results
  --api-token API_TOKEN, -a API_TOKEN
                        DigitalOcean API Token
```

'''

# (c) 2013, Evan Wies <evan@neomantra.net>
#
# Inspired by the EC2 inventory plugin:
# https://github.com/ansible/ansible/blob/devel/plugins/inventory/ec2.py
#
# This file is part of Ansible,
#
# Ansible 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.
#
# Ansible 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 Ansible.  If not, see <http://www.gnu.org/licenses/>.

######################################################################

import os
import sys
import re
import argparse
from time import time
import ConfigParser

try:
    import json
except ImportError:
    import simplejson as json

try:
    from dopy.manager import DoError, DoManager
except ImportError, e:
    print "failed=True msg='`dopy` library required for this script'"
    sys.exit(1)

class DigitalOceanInventory(object):

    ###########################################################################
    # Main execution path
    ###########################################################################

    def __init__(self):
        ''' Main execution path '''

        # DigitalOceanInventory data
        self.data = {}      # All DigitalOcean data
        self.inventory = {} # Ansible Inventory

        # Read settings, environment variables, and CLI arguments
        self.read_settings()
        self.read_environment()
        self.read_cli_args()

        # Verify credentials were set
        if not hasattr(self, 'api_token'):
            print '''Could not find values for DigitalOcean api_token.
They must be specified via either ini file, command line argument (--api-token),
or environment variables (DO_API_TOKEN)'''
            sys.exit(-1)

        # env command, show DigitalOcean credentials
        if self.args.env:
            print "DO_API_TOKEN=%s" % self.api_token
            sys.exit(0)

        self.manager = DoManager(None, self.api_token, api_version=2)

        # Pick the json_data to print based on the CLI command
        if self.args.droplets:
            json_data = self.load_from_digital_ocean('droplets')
        elif self.args.regions:
            json_data = self.load_from_digital_ocean('regions')
        elif self.args.images:
            json_data = self.load_from_digital_ocean('images')
        elif self.args.sizes:
            json_data = self.load_from_digital_ocean('sizes')
        elif self.args.ssh_keys:
            json_data = self.load_from_digital_ocean('ssh_keys')
        elif self.args.domains:
            json_data = self.load_from_digital_ocean('domains')
        elif self.args.all:
            json_data = self.load_from_digital_ocean()
        elif self.args.host:
            json_data = self.load_droplet_variables_for_host()
        else:    # '--list' this is last to make it default
            self.data = self.load_from_digital_ocean('droplets')
            self.build_inventory()
            json_data = self.inventory

        if self.args.pretty:
            print json.dumps(json_data, sort_keys=True, indent=2)
        else:
            print json.dumps(json_data)
        # That's all she wrote...

    ###########################################################################
    # Script configuration
    ###########################################################################

    def read_settings(self):
        ''' Reads the settings from the digital_ocean.ini file '''
        config = ConfigParser.SafeConfigParser()
        config.read(os.path.dirname(os.path.realpath(__file__)) + '/digital_ocean.ini')

        # Credentials
        if config.has_option('digital_ocean', 'api_token'):
            self.api_token = config.get('digital_ocean', 'api_token')

        # Cache related
        if config.has_option('digital_ocean', 'cache_path'):
            self.cache_path = config.get('digital_ocean', 'cache_path')
        if config.has_option('digital_ocean', 'cache_max_age'):
            self.cache_max_age = config.getint('digital_ocean', 'cache_max_age')

    def read_environment(self):
        ''' Reads the settings from environment variables '''
        # Setup credentials
        if os.getenv("DO_API_TOKEN"):
            self.api_token = os.getenv("DO_API_TOKEN")
        if os.getenv("DO_API_KEY"):
            self.api_token = os.getenv("DO_API_KEY")

    def read_cli_args(self):
        ''' Command line argument processing '''
        parser = argparse.ArgumentParser(description='Produce an Ansible Inventory file based on DigitalOcean credentials')

        parser.add_argument('--list', action='store_true', help='List all active Droplets as Ansible inventory (default: True)')
        parser.add_argument('--host', action='store', help='Get all Ansible inventory variables about a specific Droplet')

        parser.add_argument('--all', action='store_true', help='List all DigitalOcean information as JSON')
        parser.add_argument('--droplets','-d', action='store_true', help='List Droplets as JSON')
        parser.add_argument('--regions', action='store_true', help='List Regions as JSON')
        parser.add_argument('--images', action='store_true', help='List Images as JSON')
        parser.add_argument('--sizes', action='store_true', help='List Sizes as JSON')
        parser.add_argument('--ssh-keys', action='store_true', help='List SSH keys as JSON')
        parser.add_argument('--domains', action='store_true',help='List Domains as JSON')

        parser.add_argument('--pretty','-p', action='store_true', help='Pretty-print results')

        parser.add_argument('--env','-e', action='store_true', help='Display DO_API_TOKEN')
        parser.add_argument('--api-token','-a', action='store', help='DigitalOcean API Token')

        self.args = parser.parse_args()

        if self.args.api_token:
            self.api_token = self.args.api_token

        # Make --list default if none of the other commands are specified
        if (not self.args.droplets and not self.args.regions and
                not self.args.images and not self.args.sizes and
                not self.args.ssh_keys and not self.args.domains and
                not self.args.all and not self.args.host):
            self.args.list = True

    ###########################################################################
    # Data Management
    ###########################################################################

    def load_from_digital_ocean(self, resource=None):
        '''Get JSON from DigitalOcean API'''
        json_data = {}
        if resource == 'droplets' or resource is None:
            json_data['droplets'] = self.manager.all_active_droplets()
        if resource == 'regions' or resource is None:
            json_data['regions'] = self.manager.all_regions()
        if resource == 'images' or resource is None:
            json_data['images'] = self.manager.all_images(filter=None)
        if resource == 'sizes' or resource is None:
            json_data['sizes'] = self.manager.sizes()
        if resource == 'ssh_keys' or resource is None:
            json_data['ssh_keys'] = self.manager.all_ssh_keys()
        if resource == 'domains' or resource is None:
            json_data['domains'] = self.manager.all_domains()
        return json_data

    def build_inventory(self):
        '''Build Ansible inventory of droplets'''
        self.inventory = {}

        # add all droplets by id and name
        for droplet in self.data['droplets']:
            dest = droplet['ip_address']

            self.inventory[droplet['id']] = [dest]
            self.push(self.inventory, droplet['name'], dest)
            self.push(self.inventory, 'region_' + droplet['region']['slug'], dest)
            self.push(self.inventory, 'image_' + str(droplet['image']['id']), dest)
            self.push(self.inventory, 'size_' + droplet['size']['slug'], dest)

            image_slug = droplet['image']['slug']
            if image_slug:
                self.push(self.inventory, 'image_' + self.to_safe(image_slug), dest)
            else:
                image_name = droplet['image']['name']
                if image_name:
                    self.push(self.inventory, 'image_' + self.to_safe(image_name), dest)

            self.push(self.inventory, 'distro_' + self.to_safe(droplet['image']['distribution']), dest)
            self.push(self.inventory, 'status_' + droplet['status'], dest)

    def load_droplet_variables_for_host(self):
        '''Generate a JSON response to a --host call'''
        host = int(self.args.host)

        return self.manager.show_droplet(host)

    ###########################################################################
    # Utilities
    ###########################################################################

    def push(self, my_dict, key, element):
        ''' Pushed an element onto an array that may not have been defined in the dict '''
        if key in my_dict:
            my_dict[key].append(element)
        else:
            my_dict[key] = [element]

    def to_safe(self, word):
        ''' Converts 'bad' characters in a string to underscores so they can be used as Ansible groups '''
        return re.sub("[^A-Za-z0-9\-\.]", "_", word)

###########################################################################
# Run the script
DigitalOceanInventory()
<|endoftext|>


<commit_before>from datetime import datetime
from io import BytesIO
import json
import logging
import uuid
import pickle
import os

from django.conf import settings
from django.contrib.auth import get_user_model
from django.core.cache import cache
from django.core.files.storage import FileSystemStorage
from django.http import JsonResponse, HttpResponseServerError
from django.http import HttpResponse
from django.views.generic import FormView
from django.views.generic.edit import ProcessFormView
from django.views.generic import TemplateView
from django.utils import timezone
from django.contrib.auth import login

from paypal.standard.forms import PayPalPaymentsForm

from disease.files_utils import process_filename
from disease.files_utils import get_genome_dirpath
from disease.files_utils import get_genome_filepath
from .forms import UploadGenomeForm
from .models import AnalyzeDataOrder
from .tasks import recompute_genome_files

log = logging.getLogger(__name__)

storage = FileSystemStorage()

def upload_progress(request):
    """
    Return JSON object with information about the progress of an upload.
    """
    progress_id = ''
    if 'X-Progress-ID' in request.GET:
        progress_id = request.GET['X-Progress-ID']
    elif 'X-Progress-ID' in request.META:
        progress_id = request.META['X-Progress-ID']
    if progress_id:
        cache_key = "%s_%s" % (request.META['REMOTE_ADDR'], progress_id)
        data = cache.get(cache_key)
        log.debug('PID: %s, Upload progress cache %s',os.getpid(),  data)
        if data is None:
            data = {'length': 1, 'uploaded': 1}
        return JsonResponse(data)
    else:
        return HttpResponseServerError('Server Error: You must provide X-Progress-ID header or query param.')

class GenomeFilePathMixin(object):

    def process_filename(self, filename, filename_suffix=None):
        return process_filename(filename, filename_suffix)

    def get_dirpath(self, user=None):
        if user is None:
            user = self.request.user
        return get_genome_dirpath(user)

    def get_filepath(self, filename, user=None):
        if user is None:
            user = self.request.user
        return get_genome_filepath(user, filename)

class UploadGenome(GenomeFilePathMixin, FormView):
    template_name = 'upload_genome.html'
    form_class = UploadGenomeForm
    success_url = reversed('disease:genome_payment')

    def get_form_kwargs(self):
        kwargs = super().get_form_kwargs()
        kwargs['user'] = self.request.user
        return kwargs

    def get(self, request, *args, **kwargs):
        """
        Handles GET requests and instantiates a blank version of the form.
        """
        form = self.get_form()
        upload_id = uuid.uuid4()
        return self.render_to_response(self.get_context_data(form=form, upload_id=upload_id))

    def save_processed_data(self, data):
        buffer = BytesIO()
        pickle.dump(data, buffer)
        filename = self.process_filename(self.request.FILES['file'].name, filename_suffix='_processed')
        storage.save(self.get_filepath(filename), buffer)

    def form_valid(self, form):
        # save file
        # create AnalyzeFileOrder
        # raw_filepath = self.get_filepath(raw_filename)
        cd = form.cleaned_data
        email = cd.get('email', None)
        raw_file = cd.get('file', None)
        raw_filename = getattr(raw_file, 'name', None)
        user_model = get_user_model()
        if not self.request.user.is_authenticated():
            try:
                user = user_model.objects.get(email=email)
            except user_model.DoesNotExist:  # user doesn't have an account, create one
                user = user_model(email=email, username=email)
                user.save()

                # Dirty hack to allow user login by model
                user.backend = 'django.contrib.auth.backends.ModelBackend'
                login(self.request, user)

            # Dirty hack to fix some parts requiring request.user...
            self.request.user = user
        else:
            user = self.request.user

        storage.save(self.get_filepath(raw_filename, user=user), raw_file)
        analyze_order = AnalyzeDataOrder(uploaded_filename=raw_filename, user=user)

        if user.is_staff and user.is_active:
            log.info('User %s skipping payment due to staff membership', user)
            analyze_order.paid = timezone.now()
        analyze_order.save()
        recompute_genome_files.delay(user.pk, user.email)
        # table = process_genoome_data(data)
        # file_exists = os.path.isfile(os.path.join(settings.MEDIA_ROOT, self.get_filepath(self.process_filename(raw_filename, filename_suffix='_processed'))))
        # if self.request.user.is_authenticated() and not file_exists:
        #     self.save_processed_data(table)

        # ctx = self.get_context_data(form=form, table=table, analyzed=True)
        pos_data = analyze_order.posData()
        ctx = self.get_context_data(
            form=form, analyzed=True, pos_data=pos_data, bitpay_checkout_url=settings.BITPAY_API,
            paypal_form=PayPalPaymentsForm(
                initial=analyze_order.paypal_data(self.request))
            )
        return self.render_to_response(ctx)

class GenomePaymentView(TemplateView):
    template_name = 'upload_success.html'

class DisplayGenomeResult(GenomeFilePathMixin, TemplateView):
    template_name = 'display_genome_result.html'

    def get(self, request, *args, **kwargs):
        self.user = self.request.user
        if self.is_browsing_via_admin:
            self.user = get_user_model().objects.get(pk=int(self.request.GET['pk']))
        return super().get(request, *args, **kwargs)

    def get_genome_data(self):
        filename = self.process_filename(self.request.GET['file'], filename_suffix='_processed')
        filepath = self.get_filepath(filename)
        with storage.open(filepath) as f:
            data = pickle.load(f)
        return data

    @property
    def is_admin(self):  # TODO use permissions?
        return bool(self.request.user.is_staff and self.request.user.is_active)

    @property
    def is_browsing_via_admin(self):
        return bool(('pk' in self.request.GET) and self.is_admin)

    def get_filepath(self, filename):
        if self.is_browsing_via_admin:
            return get_genome_filepath(self.user, filename)
        return super().get_filepath(filename)

    def get_context_data(self, **kwargs):
        ctx = super().get_context_data(**kwargs)
        ctx['is_admin'] = is_admin = self.is_admin

        order_kwargs = dict(uploaded_filename=self.request.GET['file'], user=self.user)
        paid = False

        try:
            analyze_data_order = AnalyzeDataOrder.objects.get(**order_kwargs)
            paid = analyze_data_order.is_paid
        except AnalyzeDataOrder.DoesNotExist:
            if not self.is_browsing_via_admin:
                analyze_data_order = AnalyzeDataOrder(**order_kwargs)
                analyze_data_order.save()
                paid = analyze_data_order.is_paid

        ctx['paid'] = paid
        if paid or is_admin:
            ctx['table'] = self.get_genome_data()
        ctx['bitpay_checkout_url'] = settings.BITPAY_API
        ctx['analyze_data_order_pk'] = analyze_data_order.pk
        ctx['pos_data'] = analyze_data_order.posData()
        ctx['paypal_form'] = PayPalPaymentsForm(
            initial=analyze_data_order.paypal_data(self.request))
        return ctx

class PaymentStatusView(ProcessFormView, TemplateView):
    http_method_names = ['post']

    def post(self, request, *args, **kwargs):
        post_data = self.request.POST
        print(post_data)
        if post_data['status'] in {'paid', 'complete', 'confirmed'}:
            posData = json.loads(post_data['posData'])
            analyze_order_pk = posData['analyze_order_pk']
            user_pk = posData['user_pk']

            analyze_order = AnalyzeDataOrder.objects.get(pk=analyze_order_pk)
            analyze_order.paid = timezone.now()
            analyze_order.save()
        return HttpResponse('OK')

<commit_msg>Fix upload coupon<commit_after>from datetime import datetime
from io import BytesIO
import json
import logging
import uuid
import pickle
import os

from django.conf import settings
from django.contrib.auth import get_user_model
from django.core.cache import cache
from django.core.files.storage import FileSystemStorage
from django.http import JsonResponse, HttpResponseServerError
from django.http import HttpResponse
from django.views.generic import FormView
from django.views.generic.edit import ProcessFormView
from django.views.generic import TemplateView
from django.utils import timezone
from django.contrib.auth import login

from paypal.standard.forms import PayPalPaymentsForm

from disease.files_utils import process_filename
from disease.files_utils import get_genome_dirpath
from disease.files_utils import get_genome_filepath
from .forms import UploadGenomeForm
from .models import AnalyzeDataOrder
from .tasks import recompute_genome_files

log = logging.getLogger(__name__)

storage = FileSystemStorage()

def upload_progress(request):
    """
    Return JSON object with information about the progress of an upload.
    """
    progress_id = ''
    if 'X-Progress-ID' in request.GET:
        progress_id = request.GET['X-Progress-ID']
    elif 'X-Progress-ID' in request.META:
        progress_id = request.META['X-Progress-ID']
    if progress_id:
        cache_key = "%s_%s" % (request.META['REMOTE_ADDR'], progress_id)
        data = cache.get(cache_key)
        log.debug('PID: %s, Upload progress cache %s',os.getpid(),  data)
        if data is None:
            data = {'length': 1, 'uploaded': 1}
        return JsonResponse(data)
    else:
        return HttpResponseServerError('Server Error: You must provide X-Progress-ID header or query param.')

class GenomeFilePathMixin(object):

    def process_filename(self, filename, filename_suffix=None):
        return process_filename(filename, filename_suffix)

    def get_dirpath(self, user=None):
        if user is None:
            user = self.request.user
        return get_genome_dirpath(user)

    def get_filepath(self, filename, user=None):
        if user is None:
            user = self.request.user
        return get_genome_filepath(user, filename)

class UploadGenome(GenomeFilePathMixin, FormView):
    template_name = 'upload_genome.html'
    form_class = UploadGenomeForm
    success_url = reversed('disease:genome_payment')

    def get_form_kwargs(self):
        kwargs = super().get_form_kwargs()
        kwargs['user'] = self.request.user
        return kwargs

    def get(self, request, *args, **kwargs):
        """
        Handles GET requests and instantiates a blank version of the form.
        """
        form = self.get_form()
        upload_id = uuid.uuid4()
        return self.render_to_response(self.get_context_data(form=form, upload_id=upload_id))

    def save_processed_data(self, data):
        buffer = BytesIO()
        pickle.dump(data, buffer)
        filename = self.process_filename(self.request.FILES['file'].name, filename_suffix='_processed')
        storage.save(self.get_filepath(filename), buffer)

    def form_valid(self, form):
        # save file
        # create AnalyzeFileOrder
        # raw_filepath = self.get_filepath(raw_filename)
        cd = form.cleaned_data
        email = cd.get('email', None)
        raw_file = cd.get('file', None)
        raw_filename = getattr(raw_file, 'name', None)
        user_model = get_user_model()
        if not self.request.user.is_authenticated():
            try:
                user = user_model.objects.get(email=email)
            except user_model.DoesNotExist:  # user doesn't have an account, create one
                user = user_model(email=email, username=email)
                user.save()

                # Dirty hack to allow user login by model
                user.backend = 'django.contrib.auth.backends.ModelBackend'
                login(self.request, user)

            # Dirty hack to fix some parts requiring request.user...
            self.request.user = user
        else:
            user = self.request.user

        storage.save(self.get_filepath(raw_filename, user=user), raw_file)
        analyze_order = AnalyzeDataOrder(uploaded_filename=raw_filename, user=user)

        if user.is_staff and user.is_active:
            log.info('User %s skipping payment due to staff membership', user)
            analyze_order.paid = timezone.now()
        analyze_order.save()
        recompute_genome_files.delay(user.pk, user.email)
        # table = process_genoome_data(data)
        # file_exists = os.path.isfile(os.path.join(settings.MEDIA_ROOT, self.get_filepath(self.process_filename(raw_filename, filename_suffix='_processed'))))
        # if self.request.user.is_authenticated() and not file_exists:
        #     self.save_processed_data(table)

        # ctx = self.get_context_data(form=form, table=table, analyzed=True)
        pos_data = analyze_order.posData()
        ctx = self.get_context_data(
            form=form, analyzed=True, pos_data=pos_data, bitpay_checkout_url=settings.BITPAY_API,
            analyze_data_order_pk=analyze_order.pk,
            paypal_form=PayPalPaymentsForm(
                initial=analyze_order.paypal_data(self.request))
            )
        return self.render_to_response(ctx)

class GenomePaymentView(TemplateView):
    template_name = 'upload_success.html'

class DisplayGenomeResult(GenomeFilePathMixin, TemplateView):
    template_name = 'display_genome_result.html'

    def get(self, request, *args, **kwargs):
        self.user = self.request.user
        if self.is_browsing_via_admin:
            self.user = get_user_model().objects.get(pk=int(self.request.GET['pk']))
        return super().get(request, *args, **kwargs)

    def get_genome_data(self):
        filename = self.process_filename(self.request.GET['file'], filename_suffix='_processed')
        filepath = self.get_filepath(filename)
        with storage.open(filepath) as f:
            data = pickle.load(f)
        return data

    @property
    def is_admin(self):  # TODO use permissions?
        return bool(self.request.user.is_staff and self.request.user.is_active)

    @property
    def is_browsing_via_admin(self):
        return bool(('pk' in self.request.GET) and self.is_admin)

    def get_filepath(self, filename):
        if self.is_browsing_via_admin:
            return get_genome_filepath(self.user, filename)
        return super().get_filepath(filename)

    def get_context_data(self, **kwargs):
        ctx = super().get_context_data(**kwargs)
        ctx['is_admin'] = is_admin = self.is_admin

        order_kwargs = dict(uploaded_filename=self.request.GET['file'], user=self.user)
        paid = False

        try:
            analyze_data_order = AnalyzeDataOrder.objects.get(**order_kwargs)
            paid = analyze_data_order.is_paid
        except AnalyzeDataOrder.DoesNotExist:
            if not self.is_browsing_via_admin:
                analyze_data_order = AnalyzeDataOrder(**order_kwargs)
                analyze_data_order.save()
                paid = analyze_data_order.is_paid

        ctx['paid'] = paid
        if paid or is_admin:
            ctx['table'] = self.get_genome_data()
        ctx['bitpay_checkout_url'] = settings.BITPAY_API
        ctx['analyze_data_order_pk'] = analyze_data_order.pk
        ctx['pos_data'] = analyze_data_order.posData()
        ctx['paypal_form'] = PayPalPaymentsForm(
            initial=analyze_data_order.paypal_data(self.request))
        return ctx

class PaymentStatusView(ProcessFormView, TemplateView):
    http_method_names = ['post']

    def post(self, request, *args, **kwargs):
        post_data = self.request.POST
        print(post_data)
        if post_data['status'] in {'paid', 'complete', 'confirmed'}:
            posData = json.loads(post_data['posData'])
            analyze_order_pk = posData['analyze_order_pk']
            user_pk = posData['user_pk']

            analyze_order = AnalyzeDataOrder.objects.get(pk=analyze_order_pk)
            analyze_order.paid = timezone.now()
            analyze_order.save()
        return HttpResponse('OK')

<|endoftext|>


<commit_before>"""

This module defines the :class:`GenericFunction` class, which is the base for
the implementation of spatial functions in GeoAlchemy.  This module is also
where actual spatial functions are defined. Spatial functions supported by
GeoAlchemy are defined in this module. See :class:`GenericFunction` to know how
to create new spatial functions.

.. note::

    By convention the names of spatial functions are prefixed by ``ST_``.  This
    is to be consistent with PostGIS', which itself is based on the ``SQL-MM``
    standard.

Functions created by subclassing :class:`GenericFunction` can be called
in several ways:

* By using the ``func`` object, which is the SQLAlchemy standard way of calling
  a function. For example, without the ORM::

      select([func.ST_Area(lake_table.c.geom)])

  and with the ORM::

      Session.query(func.ST_Area(Lake.geom))

* By applying the function to a geometry column. For example, without the
  ORM::

      select([lake_table.c.geom.ST_Area()])

  and with the ORM::

      Session.query(Lake.geom.ST_Area())

* By applying the function to a :class:`geoalchemy2.elements.WKBElement`
  object (:class:`geoalchemy2.elements.WKBElement` is the type into
  which GeoAlchemy converts geometry values read from the database), or
  to a :class:`geoalchemy2.elements.WKTElement` object. For example,
  without the ORM::

      conn.scalar(lake['geom'].ST_Area())

  and with the ORM::

      session.scalar(lake.geom.ST_Area())

Reference
---------

"""

from sqlalchemy.sql import functions

from . import types

class GenericFunction(functions.GenericFunction):
    """
    The base class for GeoAlchemy functions.

    This class inherits from ``sqlalchemy.sql.functions.GenericFunction``, so
    functions defined by subclassing this class can be given a fixed return
    type. For example, functions like :class:`ST_Buffer` and
    :class:`ST_Envelope` have their ``type`` attributes set to
    :class:`geoalchemy2.types.Geometry`.

    This class allows constructs like ``Lake.geom.ST_Buffer(2)``. In that
    case the ``Function`` instance is bound to an expression (``Lake.geom``
    here), and that expression is passed to the function when the function
    is actually called.

    If you need to use a function that GeoAlchemy does not provide you will
    certainly want to subclass this class. For example, if you need the
    ``ST_TransScale`` spatial function, which isn't (currently) natively
    supported by GeoAlchemy, you will write this::

        from geoalchemy2 import Geometry
        from geoalchemy2.functions import GenericFunction

        class ST_TransScale(GenericFunction):
            name = 'ST_TransScale'
            type = Geometry
    """

    def __init__(self, *args, **kwargs):
        expr = kwargs.pop('expr', None)
        if expr is not None:
            args = (expr,) + args
        functions.GenericFunction.__init__(self, *args, **kwargs)

# Functions are classified as in the PostGIS doc.
# <http://www.postgis.org/documentation/manual-svn/reference.html>

_FUNCTIONS = [
    #
    # Geometry Accessors
    #

    ('ST_Envelope', types.Geometry,
     'Returns a geometry representing the double precision (float8) bounding'
     'box of the supplied geometry.'),

    ('ST_GeometryN', None,
     'Return the 1-based Nth geometry if the geometry is a '
     '``GEOMETRYCOLLECTION``, ``(MULTI)POINT``, ``(MULTI)LINESTRING``, '
     '``MULTICURVE`` or ``(MULTI)POLYGON``, ``POLYHEDRALSURFACE`` Otherwise, '
     'return ``None``.'),

    ('ST_GeometryType', None,
     'Return the geometry type of the ``ST_Geometry`` value.'),

    ('ST_IsValid', None,
     'Returns ``True`` if the ``ST_Geometry`` is well formed.'),

    ('ST_NPoints', None,
     'Return the number of points (vertices) in a geometry.'),

    ('ST_SRID', None,
     'Returns the spatial reference identifier for the ``ST_Geometry`` as '
     'defined in ``spatial_ref_sys`` table.'),

    ('ST_X', None,
     'Return the X coordinate of the point, or ``None`` if not available. '
     'Input must be a point.'),

    ('ST_Y', None,
     'Return the Y coordinate of the point, or ``None`` if not available. '
     'Input must be a point.'),

    #
    # Geometry Editors
    #

    ('ST_Transform', types.Geometry,
     'Returns a new geometry with its coordinates transformed to the SRID '
     'referenced by the integer parameter.'),

    #
    # Geometry Outputs
    #

    ('ST_AsBinary', None,
     'Return the Well-Known Binary (WKB) representation of the geometry/'
     'geography without SRID meta data.'),

    ('ST_AsGeoJSON', None, 'Return the geometry as a GeoJSON element.'),

    ('ST_AsGML', None, 'Return the geometry as a GML version 2 or 3 element.'),

    ('ST_AsKML', None,
     'Return the geometry as a KML element. Several variants. Default '
     'version=2, default precision=15'),

    ('ST_AsSVG', None,
     'Returns a Geometry in SVG path data given a geometry or geography '
     'object.'),

    ('ST_AsText', None,
     'Return the Well-Known Text (WKT) representation of the geometry/'
     'geography without SRID metadata.'),

    #
    # Spatial Relationships and Measurements
    #

    ('ST_Area', None,
     'Returns the area of the surface if it is a polygon or multi-polygon. '
     'For ``geometry`` type area is in SRID units. For ``geography`` area is '
     'in square meters.'),

    ('ST_Centroid', types.Geometry,
     'Returns the geometric center of a geometry.'),

    ('ST_Contains', None,
     'Returns ``True`` if and only if no points of B lie in the exterior of '
     'A, and at least one point of the interior of B lies in the interior '
     'of A.'),

    ('ST_ContainsProperly', None,
     'Returns ``True`` if B intersects the interior of A but not the boundary '
     '(or exterior). A does not contain properly itself, but does contain '
     'itself.'),

    ('ST_Covers', None,
     'Returns ``True`` if no point in Geometry B is outside Geometry A'),

    ('ST_CoveredBy', None,
     'Returns ``True`` if no point in Geometry/Geography A is outside Geometry'
     '/Geography B'),

    ('ST_Crosses', None,
     'Returns ``True`` if the supplied geometries have some, but not all, '
     'interior points in common.'),

    ('ST_Disjoint', None,
     'Returns ``True`` if the Geometries do not "spatially intersect" - if '
     'they do not share any space together.'),

    ('ST_Distance', None,
     'For geometry type Returns the 2-dimensional cartesian minimum distance '
     '(based on spatial ref) between two geometries in projected units. For '
     'geography type defaults to return spheroidal minimum distance between '
     'two geographies in meters.'),

    ('ST_Distance_Sphere', None,
     'Returns minimum distance in meters between two lon/lat geometries. Uses '
     'a spherical earth and radius of 6370986 meters. Faster than '
     '``ST_Distance_Spheroid``, but less accurate. PostGIS versions '
     'prior to 1.5 only implemented for points.'),

    ('ST_DFullyWithin', None,
     'Returns ``True`` if all of the geometries are within the specified '
     'distance of one another'),

    ('ST_DWithin', None,
     'Returns ``True`` if the geometries are within the specified distance of '
     'one another. For geometry units are in those of spatial reference and '
     'for geography units are in meters and measurement is defaulted to '
     '``use_spheroid=True`` (measure around spheroid), for faster check, '
     '``use_spheroid=False`` to measure along sphere.'),

    ('ST_Equals', None,
     'Returns ``True`` if the given geometries represent the same geometry. '
     'Directionality is ignored.'),

    ('ST_Intersects', None,
     'Returns ``True`` if the Geometries/Geography "spatially intersect in '
     '2D" - (share any portion of space) and ``False`` if they don\'t (they '
     'are Disjoint). For geography -- tolerance is 0.00001 meters (so any '
     'points that close are considered to intersect)'),

    ('ST_Length', None,
     'Returns the 2d length of the geometry if it is a linestring or '
     'multilinestring. geometry are in units of spatial reference and '
     'geography are in meters (default spheroid)'),

    ('ST_OrderingEquals', None,
     'Returns ``True`` if the given geometries represent the same geometry '
     'and points are in the same directional order.'),

    ('ST_Overlaps', None,
     'Returns ``True`` if the Geometries share space, are of the same '
     'dimension, but are not completely contained by each other.'),

    ('ST_Perimeter', None,
     'Return the length measurement of the boundary of an ST_Surface or '
     'ST_MultiSurface geometry or geography. (Polygon, Multipolygon). '
     'geometry measurement is in units of spatial reference and geography is '
     'in meters.'),

    ('ST_Project', types.Geography,
     'Returns a ``POINT`` projected from a start point using a distance in '
     'meters and bearing (azimuth) in radians.'),

    ('ST_Relate', None,
     'Returns ``True`` if this Geometry is spatially related to '
     'anotherGeometry, by testing for intersections between the Interior, '
     'Boundary and Exterior of the two geometries as specified by the values '
     'in the intersectionMatrixPattern. If no intersectionMatrixPattern is '
     'passed in, then returns the maximum intersectionMatrixPattern that '
     'relates the 2 geometries.'),

    ('ST_Touches', None,
     'Returns ``True`` if the geometries have at least one point in common, '
     'but their interiors do not intersect.'),

    ('ST_Within', None,
     'Returns ``True`` if the geometry A is completely inside geometry B'),

    #
    # Geometry Processing
    #

    ('ST_Buffer', types.Geometry,
     'For geometry: Returns a geometry that represents all points whose '
     'distance from this Geometry is less than or equal to distance. '
     'Calculations are in the Spatial Reference System of this Geometry.\n\n'
     'For geography: Uses a planar transform wrapper. Introduced in 1.5 '
     'support for different end cap and mitre settings to control shape.'),

    ('ST_Difference', types.Geometry,
     'Returns a geometry that represents that part of geometry A that does '
     'not intersect with geometry B.'),

    ('ST_Dump', types.GeometryDump,
     'Returns a set of geometry_dump (geom,path) rows, that make up a '
     'geometry g1.'),

    ('ST_DumpPoints', types.GeometryDump,
     'Returns a set of geometry_dump (geom,path) rows of all points that '
     'make up a geometry.'),

    ('ST_Intersection', types.Geometry,
     'Returns a geometry that represents the shared portion of geomA and '
     'geomB. The geography implementation does a transform to geometry to do '
     'the intersection and then transform back to WGS84.'),

    ('ST_Union', types.Geometry,
     'Returns a geometry that represents the point set union of the '
     'Geometries.'),

    ('ST_Simplify', types.Geometry,
     'Returns a "simplified" version of the given geometry using the '
     'Douglas-Peucker algorithm')

    #
    # Raster Constructors
    #

    ('ST_AsRaster', types.Raster,
     ('Converts a PostGIS geometry to a PostGIS raster.', 'RT_ST_AsRaster')),

    #
    # Raster Accessors
    #

    ('ST_Height', None,
     ('Returns the height of the raster in pixels.', 'RT_ST_Height')),

    ('ST_Width', None,
     ('Returns the width of the raster in pixels.', 'RT_ST_Width')),

    #
    # Raster Pixel Accessors and Setters
    #

    ('ST_Value', None,
     ('Returns the value of a given band in a given columnx, rowy pixel or at '
      'a particular geometric point. Band numbers start at 1 and assumed to '
      'be 1 if not specified. If ``exclude_nodata_value`` is set to '
      '``false``, then all pixels include nodata pixels are considered to '
      'intersect and return value. If ``exclude_nodata_value`` is not passed '
      'in then reads it from metadata of raster.', 'RT_ST_Value')),
]

# Iterate through _FUNCTION and create GenericFunction classes dynamically
for name, type_, doc in _FUNCTIONS:
    attributes = {'name': name}
    docs = []

    if isinstance(doc, tuple):
        docs.append(doc[0])
        docs.append('see http://postgis.net/docs/{0}.html'.format(doc[1]))
    elif doc is not None:
        docs.append(doc)
        docs.append('see http://postgis.net/docs/{0}.html'.format(name))

    if type_ is not None:
        attributes['type'] = type_

        type_str = '{0}.{1}'.format(type_.__module__, type_.__name__)
        docs.append('Return type: :class:`{0}`.'.format(type_str))

    if len(docs) != 0:
        attributes['__doc__'] = '\n\n'.join(docs)

    globals()[name] = type(name, (GenericFunction,), attributes)
<commit_msg>Fix syntax error<commit_after>"""

This module defines the :class:`GenericFunction` class, which is the base for
the implementation of spatial functions in GeoAlchemy.  This module is also
where actual spatial functions are defined. Spatial functions supported by
GeoAlchemy are defined in this module. See :class:`GenericFunction` to know how
to create new spatial functions.

.. note::

    By convention the names of spatial functions are prefixed by ``ST_``.  This
    is to be consistent with PostGIS', which itself is based on the ``SQL-MM``
    standard.

Functions created by subclassing :class:`GenericFunction` can be called
in several ways:

* By using the ``func`` object, which is the SQLAlchemy standard way of calling
  a function. For example, without the ORM::

      select([func.ST_Area(lake_table.c.geom)])

  and with the ORM::

      Session.query(func.ST_Area(Lake.geom))

* By applying the function to a geometry column. For example, without the
  ORM::

      select([lake_table.c.geom.ST_Area()])

  and with the ORM::

      Session.query(Lake.geom.ST_Area())

* By applying the function to a :class:`geoalchemy2.elements.WKBElement`
  object (:class:`geoalchemy2.elements.WKBElement` is the type into
  which GeoAlchemy converts geometry values read from the database), or
  to a :class:`geoalchemy2.elements.WKTElement` object. For example,
  without the ORM::

      conn.scalar(lake['geom'].ST_Area())

  and with the ORM::

      session.scalar(lake.geom.ST_Area())

Reference
---------

"""

from sqlalchemy.sql import functions

from . import types

class GenericFunction(functions.GenericFunction):
    """
    The base class for GeoAlchemy functions.

    This class inherits from ``sqlalchemy.sql.functions.GenericFunction``, so
    functions defined by subclassing this class can be given a fixed return
    type. For example, functions like :class:`ST_Buffer` and
    :class:`ST_Envelope` have their ``type`` attributes set to
    :class:`geoalchemy2.types.Geometry`.

    This class allows constructs like ``Lake.geom.ST_Buffer(2)``. In that
    case the ``Function`` instance is bound to an expression (``Lake.geom``
    here), and that expression is passed to the function when the function
    is actually called.

    If you need to use a function that GeoAlchemy does not provide you will
    certainly want to subclass this class. For example, if you need the
    ``ST_TransScale`` spatial function, which isn't (currently) natively
    supported by GeoAlchemy, you will write this::

        from geoalchemy2 import Geometry
        from geoalchemy2.functions import GenericFunction

        class ST_TransScale(GenericFunction):
            name = 'ST_TransScale'
            type = Geometry
    """

    def __init__(self, *args, **kwargs):
        expr = kwargs.pop('expr', None)
        if expr is not None:
            args = (expr,) + args
        functions.GenericFunction.__init__(self, *args, **kwargs)

# Functions are classified as in the PostGIS doc.
# <http://www.postgis.org/documentation/manual-svn/reference.html>

_FUNCTIONS = [
    #
    # Geometry Accessors
    #

    ('ST_Envelope', types.Geometry,
     'Returns a geometry representing the double precision (float8) bounding'
     'box of the supplied geometry.'),

    ('ST_GeometryN', None,
     'Return the 1-based Nth geometry if the geometry is a '
     '``GEOMETRYCOLLECTION``, ``(MULTI)POINT``, ``(MULTI)LINESTRING``, '
     '``MULTICURVE`` or ``(MULTI)POLYGON``, ``POLYHEDRALSURFACE`` Otherwise, '
     'return ``None``.'),

    ('ST_GeometryType', None,
     'Return the geometry type of the ``ST_Geometry`` value.'),

    ('ST_IsValid', None,
     'Returns ``True`` if the ``ST_Geometry`` is well formed.'),

    ('ST_NPoints', None,
     'Return the number of points (vertices) in a geometry.'),

    ('ST_SRID', None,
     'Returns the spatial reference identifier for the ``ST_Geometry`` as '
     'defined in ``spatial_ref_sys`` table.'),

    ('ST_X', None,
     'Return the X coordinate of the point, or ``None`` if not available. '
     'Input must be a point.'),

    ('ST_Y', None,
     'Return the Y coordinate of the point, or ``None`` if not available. '
     'Input must be a point.'),

    #
    # Geometry Editors
    #

    ('ST_Transform', types.Geometry,
     'Returns a new geometry with its coordinates transformed to the SRID '
     'referenced by the integer parameter.'),

    #
    # Geometry Outputs
    #

    ('ST_AsBinary', None,
     'Return the Well-Known Binary (WKB) representation of the geometry/'
     'geography without SRID meta data.'),

    ('ST_AsGeoJSON', None, 'Return the geometry as a GeoJSON element.'),

    ('ST_AsGML', None, 'Return the geometry as a GML version 2 or 3 element.'),

    ('ST_AsKML', None,
     'Return the geometry as a KML element. Several variants. Default '
     'version=2, default precision=15'),

    ('ST_AsSVG', None,
     'Returns a Geometry in SVG path data given a geometry or geography '
     'object.'),

    ('ST_AsText', None,
     'Return the Well-Known Text (WKT) representation of the geometry/'
     'geography without SRID metadata.'),

    #
    # Spatial Relationships and Measurements
    #

    ('ST_Area', None,
     'Returns the area of the surface if it is a polygon or multi-polygon. '
     'For ``geometry`` type area is in SRID units. For ``geography`` area is '
     'in square meters.'),

    ('ST_Centroid', types.Geometry,
     'Returns the geometric center of a geometry.'),

    ('ST_Contains', None,
     'Returns ``True`` if and only if no points of B lie in the exterior of '
     'A, and at least one point of the interior of B lies in the interior '
     'of A.'),

    ('ST_ContainsProperly', None,
     'Returns ``True`` if B intersects the interior of A but not the boundary '
     '(or exterior). A does not contain properly itself, but does contain '
     'itself.'),

    ('ST_Covers', None,
     'Returns ``True`` if no point in Geometry B is outside Geometry A'),

    ('ST_CoveredBy', None,
     'Returns ``True`` if no point in Geometry/Geography A is outside Geometry'
     '/Geography B'),

    ('ST_Crosses', None,
     'Returns ``True`` if the supplied geometries have some, but not all, '
     'interior points in common.'),

    ('ST_Disjoint', None,
     'Returns ``True`` if the Geometries do not "spatially intersect" - if '
     'they do not share any space together.'),

    ('ST_Distance', None,
     'For geometry type Returns the 2-dimensional cartesian minimum distance '
     '(based on spatial ref) between two geometries in projected units. For '
     'geography type defaults to return spheroidal minimum distance between '
     'two geographies in meters.'),

    ('ST_Distance_Sphere', None,
     'Returns minimum distance in meters between two lon/lat geometries. Uses '
     'a spherical earth and radius of 6370986 meters. Faster than '
     '``ST_Distance_Spheroid``, but less accurate. PostGIS versions '
     'prior to 1.5 only implemented for points.'),

    ('ST_DFullyWithin', None,
     'Returns ``True`` if all of the geometries are within the specified '
     'distance of one another'),

    ('ST_DWithin', None,
     'Returns ``True`` if the geometries are within the specified distance of '
     'one another. For geometry units are in those of spatial reference and '
     'for geography units are in meters and measurement is defaulted to '
     '``use_spheroid=True`` (measure around spheroid), for faster check, '
     '``use_spheroid=False`` to measure along sphere.'),

    ('ST_Equals', None,
     'Returns ``True`` if the given geometries represent the same geometry. '
     'Directionality is ignored.'),

    ('ST_Intersects', None,
     'Returns ``True`` if the Geometries/Geography "spatially intersect in '
     '2D" - (share any portion of space) and ``False`` if they don\'t (they '
     'are Disjoint). For geography -- tolerance is 0.00001 meters (so any '
     'points that close are considered to intersect)'),

    ('ST_Length', None,
     'Returns the 2d length of the geometry if it is a linestring or '
     'multilinestring. geometry are in units of spatial reference and '
     'geography are in meters (default spheroid)'),

    ('ST_OrderingEquals', None,
     'Returns ``True`` if the given geometries represent the same geometry '
     'and points are in the same directional order.'),

    ('ST_Overlaps', None,
     'Returns ``True`` if the Geometries share space, are of the same '
     'dimension, but are not completely contained by each other.'),

    ('ST_Perimeter', None,
     'Return the length measurement of the boundary of an ST_Surface or '
     'ST_MultiSurface geometry or geography. (Polygon, Multipolygon). '
     'geometry measurement is in units of spatial reference and geography is '
     'in meters.'),

    ('ST_Project', types.Geography,
     'Returns a ``POINT`` projected from a start point using a distance in '
     'meters and bearing (azimuth) in radians.'),

    ('ST_Relate', None,
     'Returns ``True`` if this Geometry is spatially related to '
     'anotherGeometry, by testing for intersections between the Interior, '
     'Boundary and Exterior of the two geometries as specified by the values '
     'in the intersectionMatrixPattern. If no intersectionMatrixPattern is '
     'passed in, then returns the maximum intersectionMatrixPattern that '
     'relates the 2 geometries.'),

    ('ST_Touches', None,
     'Returns ``True`` if the geometries have at least one point in common, '
     'but their interiors do not intersect.'),

    ('ST_Within', None,
     'Returns ``True`` if the geometry A is completely inside geometry B'),

    #
    # Geometry Processing
    #

    ('ST_Buffer', types.Geometry,
     'For geometry: Returns a geometry that represents all points whose '
     'distance from this Geometry is less than or equal to distance. '
     'Calculations are in the Spatial Reference System of this Geometry.\n\n'
     'For geography: Uses a planar transform wrapper. Introduced in 1.5 '
     'support for different end cap and mitre settings to control shape.'),

    ('ST_Difference', types.Geometry,
     'Returns a geometry that represents that part of geometry A that does '
     'not intersect with geometry B.'),

    ('ST_Dump', types.GeometryDump,
     'Returns a set of geometry_dump (geom,path) rows, that make up a '
     'geometry g1.'),

    ('ST_DumpPoints', types.GeometryDump,
     'Returns a set of geometry_dump (geom,path) rows of all points that '
     'make up a geometry.'),

    ('ST_Intersection', types.Geometry,
     'Returns a geometry that represents the shared portion of geomA and '
     'geomB. The geography implementation does a transform to geometry to do '
     'the intersection and then transform back to WGS84.'),

    ('ST_Union', types.Geometry,
     'Returns a geometry that represents the point set union of the '
     'Geometries.'),

    ('ST_Simplify', types.Geometry,
     'Returns a "simplified" version of the given geometry using the '
     'Douglas-Peucker algorithm'),

    #
    # Raster Constructors
    #

    ('ST_AsRaster', types.Raster,
     ('Converts a PostGIS geometry to a PostGIS raster.', 'RT_ST_AsRaster')),

    #
    # Raster Accessors
    #

    ('ST_Height', None,
     ('Returns the height of the raster in pixels.', 'RT_ST_Height')),

    ('ST_Width', None,
     ('Returns the width of the raster in pixels.', 'RT_ST_Width')),

    #
    # Raster Pixel Accessors and Setters
    #

    ('ST_Value', None,
     ('Returns the value of a given band in a given columnx, rowy pixel or at '
      'a particular geometric point. Band numbers start at 1 and assumed to '
      'be 1 if not specified. If ``exclude_nodata_value`` is set to '
      '``false``, then all pixels include nodata pixels are considered to '
      'intersect and return value. If ``exclude_nodata_value`` is not passed '
      'in then reads it from metadata of raster.', 'RT_ST_Value')),
]

# Iterate through _FUNCTION and create GenericFunction classes dynamically
for name, type_, doc in _FUNCTIONS:
    attributes = {'name': name}
    docs = []

    if isinstance(doc, tuple):
        docs.append(doc[0])
        docs.append('see http://postgis.net/docs/{0}.html'.format(doc[1]))
    elif doc is not None:
        docs.append(doc)
        docs.append('see http://postgis.net/docs/{0}.html'.format(name))

    if type_ is not None:
        attributes['type'] = type_

        type_str = '{0}.{1}'.format(type_.__module__, type_.__name__)
        docs.append('Return type: :class:`{0}`.'.format(type_str))

    if len(docs) != 0:
        attributes['__doc__'] = '\n\n'.join(docs)

    globals()[name] = type(name, (GenericFunction,), attributes)
<|endoftext|>


<commit_before>import urllib2
import lxml.html
import numpy
import scipy
import scipy.misc
import scipy.cluster
import urlparse
import struct
import operator
import gzip
import datetime
import requests
from PIL import BmpImagePlugin, PngImagePlugin, Image
from socket import error as SocketError
from boto.s3.key import Key
from StringIO import StringIO
from django.conf import settings
from apps.rss_feeds.models import MFeedPage, MFeedIcon
from utils import log as logging
from utils.feed_functions import timelimit, TimeoutError
from OpenSSL.SSL import Error as OpenSSLError
from pyasn1.error import PyAsn1Error
from requests.packages.urllib3.exceptions import LocationParseError

class IconImporter(object):

    def __init__(self, feed, page_data=None, force=False):
        self.feed = feed
        self.force = force
        self.page_data = page_data
        self.feed_icon = MFeedIcon.get_feed(feed_id=self.feed.pk)

    def save(self):
        if not self.force and self.feed.favicon_not_found:
            # print 'Not found, skipping...'
            return
        if (not self.force and
            not self.feed.favicon_not_found and
            self.feed_icon.icon_url and
            self.feed.s3_icon):
            # print 'Found, but skipping...'
            return
        image, image_file, icon_url = self.fetch_image_from_page_data()
        if not image:
            image, image_file, icon_url = self.fetch_image_from_path(force=self.force)

        if image:
            image     = self.normalize_image(image)
            try:
                color     = self.determine_dominant_color_in_image(image)
            except IndexError:
                return
            try:
                image_str = self.string_from_image(image)
            except TypeError:
                return

            if len(image_str) > 500000:
                image = None
            if (image and
                (self.force or
                 self.feed_icon.color != color or
                 self.feed_icon.data != image_str or
                 self.feed_icon.icon_url != icon_url or
                 self.feed_icon.not_found or
                 (settings.BACKED_BY_AWS.get('icons_on_s3') and not self.feed.s3_icon))):
                self.feed_icon.data      = image_str
                self.feed_icon.icon_url  = icon_url
                self.feed_icon.color     = color
                self.feed_icon.not_found = False
                self.feed_icon.save()
                if settings.BACKED_BY_AWS.get('icons_on_s3'):
                    self.save_to_s3(image_str)
            self.feed.favicon_color     = color
            self.feed.favicon_not_found = False

        if not image:
            self.feed_icon.not_found = True
            self.feed.favicon_not_found = True

        self.feed.save()
        return not self.feed.favicon_not_found

    def save_to_s3(self, image_str):
        expires = datetime.datetime.now() + datetime.timedelta(days=60)
        expires = expires.strftime("%a, %d %b %Y %H:%M:%S GMT")
        k = Key(settings.S3_ICONS_BUCKET)
        k.key = self.feed.s3_icons_key
        k.set_metadata('Content-Type', 'image/png')
        k.set_metadata('Expires', expires)
        k.set_contents_from_string(image_str.decode('base64'))
        k.set_acl('public-read')

        self.feed.s3_icon = True

    def load_icon(self, image_file, index=None):
        '''
        DEPRECATED

        Load Windows ICO image.

        See http://en.wikipedia.org/w/index.php?oldid=264332061 for file format
        description.

        Cribbed and modified from http://djangosnippets.org/snippets/1287/
        '''
        try:
            image_file.seek(0)
            header = struct.unpack('<3H', image_file.read(6))
        except Exception, e:
            return

        # Check magic
        if header[:2] != (0, 1):
            return

        # Collect icon directories
        directories = []
        for i in xrange(header[2]):
            directory = list(struct.unpack('<4B2H2I', image_file.read(16)))
            for j in xrange(3):
                if not directory[j]:
                    directory[j] = 256

            directories.append(directory)

        if index is None:
            # Select best icon
            directory = max(directories, key=operator.itemgetter(slice(0, 3)))
        else:
            directory = directories[index]

        # Seek to the bitmap data
        image_file.seek(directory[7])

        prefix = image_file.read(16)
        image_file.seek(-16, 1)

        if PngImagePlugin._accept(prefix):
            # Windows Vista icon with PNG inside
            try:
                image = PngImagePlugin.PngImageFile(image_file)
            except IOError:
                return
        else:
            # Load XOR bitmap
            try:
                image = BmpImagePlugin.DibImageFile(image_file)
            except IOError:
                return
            if image.mode == 'RGBA':
                # Windows XP 32-bit color depth icon without AND bitmap
                pass
            else:
                # Patch up the bitmap height
                image.size = image.size[0], image.size[1] >> 1
                d, e, o, a = image.tile[0]
                image.tile[0] = d, (0, 0) + image.size, o, a

                # Calculate AND bitmap dimensions. See
                # http://en.wikipedia.org/w/index.php?oldid=264236948#Pixel_storage
                # for description
                offset = o + a[1] * image.size[1]
                stride = ((image.size[0] + 31) >> 5) << 2
                size = stride * image.size[1]

                # Load AND bitmap
                image_file.seek(offset)
                string = image_file.read(size)
                mask = Image.frombytes('1', image.size, string, 'raw',
                                       ('1;I', stride, -1))

                image = image.convert('RGBA')
                image.putalpha(mask)

        return image

    def fetch_image_from_page_data(self):
        image = None
        image_file = None
        if self.page_data:
            content = self.page_data
        elif settings.BACKED_BY_AWS.get('pages_on_s3') and self.feed.s3_page:
            key = settings.S3_PAGES_BUCKET.get_key(self.feed.s3_pages_key)
            compressed_content = key.get_contents_as_string()
            stream = StringIO(compressed_content)
            gz = gzip.GzipFile(fileobj=stream)
            try:
                content = gz.read()
            except IOError:
                content = None
        else:
            content = MFeedPage.get_data(feed_id=self.feed.pk)
        url = self._url_from_html(content)
        if not url:
            try:
                content = requests.get(self.feed.feed_link).content
                url = self._url_from_html(content)
            except (AttributeError, SocketError, requests.ConnectionError,
                    requests.models.MissingSchema, requests.sessions.InvalidSchema,
                    LocationParseError, OpenSSLError, PyAsn1Error), e:
                logging.debug(" ---> ~SN~FRFailed~FY to fetch ~FGfeed icon~FY: %s" % e)
        if url:
            image, image_file = self.get_image_from_url(url)
        return image, image_file, url

    def fetch_image_from_path(self, path='favicon.ico', force=False):
        image = None
        url = None

        if not force:
            url = self.feed_icon.icon_url
        if not url and self.feed.feed_link and len(self.feed.feed_link) > 6:
            url = urlparse.urljoin(self.feed.feed_link, 'favicon.ico')
        if not url: return None, None, None

        image, image_file = self.get_image_from_url(url)
        if not image:
            url = urlparse.urljoin(self.feed.feed_link, '/favicon.ico')
            image, image_file = self.get_image_from_url(url)
        # print 'Found: %s - %s' % (url, image)
        return image, image_file, url

    def get_image_from_url(self, url):
        # print 'Requesting: %s' % url
        if not url:
            return None, None

        @timelimit(30)
        def _1(url):
            headers = {
                'User-Agent': 'NewsBlur Favicon Fetcher - %s subscriber%s - %s '
                              '(Mozilla/5.0 (Macintosh; Intel Mac OS X 10_7_1) '
                              'AppleWebKit/534.48.3 (KHTML, like Gecko) Version/5.1 '
                              'Safari/534.48.3)' % (
                    self.feed.num_subscribers,
                    's' if self.feed.num_subscribers != 1 else '',
                    self.feed.permalink,
                ),
                'Connection': 'close',
            }
            try:
                request = urllib2.Request(url, headers=headers)
                icon = urllib2.urlopen(request).read()
            except Exception:
                return None
            return icon
        try:
            icon = _1(url)
        except TimeoutError:
            return None, None

        try:
            icon_file = StringIO(icon)
            image = Image.open(icon_file)
        except (IOError, ValueError):
            return None, None

        return image, icon_file

    def _url_from_html(self, content):
        url = None
        if not content: return url
        try:
            if isinstance(content, unicode):
                content = content.encode('utf-8')
            icon_path = lxml.html.fromstring(content).xpath(
                '//link[@rel="icon" or @rel="shortcut icon"]/@href'
            )
        except (lxml.etree.ParserError, TypeError):
            return url

        if icon_path:
            if str(icon_path[0]).startswith('http'):
                url = icon_path[0]
            else:
                url = urlparse.urljoin(self.feed.feed_link, icon_path[0])
        return url

    def normalize_image(self, image):
        # if image.size != (16, 16):
        #     image = image.resize((16, 16), Image.BICUBIC)
        if image.mode != 'RGBA':
            try:
                image = image.convert('RGBA')
            except IOError:
                pass

        return image

    def determine_dominant_color_in_image(self, image):
        NUM_CLUSTERS = 5

        # Convert image into array of values for each point.
        if image.mode == '1':
            image.convert('L')
        ar = numpy.array(image)
        # ar = scipy.misc.fromimage(image)
        shape = ar.shape

        # Reshape array of values to merge color bands. [[R], [G], [B], [A]] => [R, G, B, A]
        if len(shape) > 2:
            ar = ar.reshape(scipy.product(shape[:2]), shape[2])

        # Get NUM_CLUSTERS worth of centroids.
        codes, _ = scipy.cluster.vq.kmeans(ar, NUM_CLUSTERS)

        # Pare centroids, removing blacks and whites and shades of really dark and really light.
        original_codes = codes
        for low, hi in [(60, 200), (35, 230), (10, 250)]:
            codes = scipy.array([code for code in codes
                                 if not ((code[0] < low and code[1] < low and code[2] < low) or
                                         (code[0] > hi and code[1] > hi and code[2] > hi))])
            if not len(codes): codes = original_codes
            else: break

        # Assign codes (vector quantization). Each vector is compared to the centroids
        # and assigned the nearest one.
        vecs, _ = scipy.cluster.vq.vq(ar, codes)

        # Count occurences of each clustered vector.
        counts, bins = scipy.histogram(vecs, len(codes))

        # Show colors for each code in its hex value.
        # colors = [''.join(chr(c) for c in code).encode('hex') for code in codes]
        # total = scipy.sum(counts)
        # print dict(zip(colors, [count/float(total) for count in counts]))

        # Find the most frequent color, based on the counts.
        index_max = scipy.argmax(counts)
        peak = codes[index_max]
        color = ''.join(chr(c) for c in peak).encode('hex')

        return color[:6]

    def string_from_image(self, image):
        output = StringIO()
        image.save(output, 'png', quality=95)
        contents = output.getvalue()
        output.close()
        return contents.encode('base64')
    <commit_msg>Handling requests issue with too many redirects.<commit_after>import urllib2
import lxml.html
import numpy
import scipy
import scipy.misc
import scipy.cluster
import urlparse
import struct
import operator
import gzip
import datetime
import requests
from PIL import BmpImagePlugin, PngImagePlugin, Image
from socket import error as SocketError
from boto.s3.key import Key
from StringIO import StringIO
from django.conf import settings
from apps.rss_feeds.models import MFeedPage, MFeedIcon
from utils import log as logging
from utils.feed_functions import timelimit, TimeoutError
from OpenSSL.SSL import Error as OpenSSLError
from pyasn1.error import PyAsn1Error
from requests.packages.urllib3.exceptions import LocationParseError

class IconImporter(object):

    def __init__(self, feed, page_data=None, force=False):
        self.feed = feed
        self.force = force
        self.page_data = page_data
        self.feed_icon = MFeedIcon.get_feed(feed_id=self.feed.pk)

    def save(self):
        if not self.force and self.feed.favicon_not_found:
            # print 'Not found, skipping...'
            return
        if (not self.force and
            not self.feed.favicon_not_found and
            self.feed_icon.icon_url and
            self.feed.s3_icon):
            # print 'Found, but skipping...'
            return
        image, image_file, icon_url = self.fetch_image_from_page_data()
        if not image:
            image, image_file, icon_url = self.fetch_image_from_path(force=self.force)

        if image:
            image     = self.normalize_image(image)
            try:
                color     = self.determine_dominant_color_in_image(image)
            except IndexError:
                return
            try:
                image_str = self.string_from_image(image)
            except TypeError:
                return

            if len(image_str) > 500000:
                image = None
            if (image and
                (self.force or
                 self.feed_icon.color != color or
                 self.feed_icon.data != image_str or
                 self.feed_icon.icon_url != icon_url or
                 self.feed_icon.not_found or
                 (settings.BACKED_BY_AWS.get('icons_on_s3') and not self.feed.s3_icon))):
                self.feed_icon.data      = image_str
                self.feed_icon.icon_url  = icon_url
                self.feed_icon.color     = color
                self.feed_icon.not_found = False
                self.feed_icon.save()
                if settings.BACKED_BY_AWS.get('icons_on_s3'):
                    self.save_to_s3(image_str)
            self.feed.favicon_color     = color
            self.feed.favicon_not_found = False

        if not image:
            self.feed_icon.not_found = True
            self.feed.favicon_not_found = True

        self.feed.save()
        return not self.feed.favicon_not_found

    def save_to_s3(self, image_str):
        expires = datetime.datetime.now() + datetime.timedelta(days=60)
        expires = expires.strftime("%a, %d %b %Y %H:%M:%S GMT")
        k = Key(settings.S3_ICONS_BUCKET)
        k.key = self.feed.s3_icons_key
        k.set_metadata('Content-Type', 'image/png')
        k.set_metadata('Expires', expires)
        k.set_contents_from_string(image_str.decode('base64'))
        k.set_acl('public-read')

        self.feed.s3_icon = True

    def load_icon(self, image_file, index=None):
        '''
        DEPRECATED

        Load Windows ICO image.

        See http://en.wikipedia.org/w/index.php?oldid=264332061 for file format
        description.

        Cribbed and modified from http://djangosnippets.org/snippets/1287/
        '''
        try:
            image_file.seek(0)
            header = struct.unpack('<3H', image_file.read(6))
        except Exception, e:
            return

        # Check magic
        if header[:2] != (0, 1):
            return

        # Collect icon directories
        directories = []
        for i in xrange(header[2]):
            directory = list(struct.unpack('<4B2H2I', image_file.read(16)))
            for j in xrange(3):
                if not directory[j]:
                    directory[j] = 256

            directories.append(directory)

        if index is None:
            # Select best icon
            directory = max(directories, key=operator.itemgetter(slice(0, 3)))
        else:
            directory = directories[index]

        # Seek to the bitmap data
        image_file.seek(directory[7])

        prefix = image_file.read(16)
        image_file.seek(-16, 1)

        if PngImagePlugin._accept(prefix):
            # Windows Vista icon with PNG inside
            try:
                image = PngImagePlugin.PngImageFile(image_file)
            except IOError:
                return
        else:
            # Load XOR bitmap
            try:
                image = BmpImagePlugin.DibImageFile(image_file)
            except IOError:
                return
            if image.mode == 'RGBA':
                # Windows XP 32-bit color depth icon without AND bitmap
                pass
            else:
                # Patch up the bitmap height
                image.size = image.size[0], image.size[1] >> 1
                d, e, o, a = image.tile[0]
                image.tile[0] = d, (0, 0) + image.size, o, a

                # Calculate AND bitmap dimensions. See
                # http://en.wikipedia.org/w/index.php?oldid=264236948#Pixel_storage
                # for description
                offset = o + a[1] * image.size[1]
                stride = ((image.size[0] + 31) >> 5) << 2
                size = stride * image.size[1]

                # Load AND bitmap
                image_file.seek(offset)
                string = image_file.read(size)
                mask = Image.frombytes('1', image.size, string, 'raw',
                                       ('1;I', stride, -1))

                image = image.convert('RGBA')
                image.putalpha(mask)

        return image

    def fetch_image_from_page_data(self):
        image = None
        image_file = None
        if self.page_data:
            content = self.page_data
        elif settings.BACKED_BY_AWS.get('pages_on_s3') and self.feed.s3_page:
            key = settings.S3_PAGES_BUCKET.get_key(self.feed.s3_pages_key)
            compressed_content = key.get_contents_as_string()
            stream = StringIO(compressed_content)
            gz = gzip.GzipFile(fileobj=stream)
            try:
                content = gz.read()
            except IOError:
                content = None
        else:
            content = MFeedPage.get_data(feed_id=self.feed.pk)
        url = self._url_from_html(content)
        if not url:
            try:
                content = requests.get(self.feed.feed_link).content
                url = self._url_from_html(content)
            except (AttributeError, SocketError, requests.ConnectionError,
                    requests.models.MissingSchema, requests.sessions.InvalidSchema,
                    requests.sessions.TooManyRedirects,
                    LocationParseError, OpenSSLError, PyAsn1Error), e:
                logging.debug(" ---> ~SN~FRFailed~FY to fetch ~FGfeed icon~FY: %s" % e)
        if url:
            image, image_file = self.get_image_from_url(url)
        return image, image_file, url

    def fetch_image_from_path(self, path='favicon.ico', force=False):
        image = None
        url = None

        if not force:
            url = self.feed_icon.icon_url
        if not url and self.feed.feed_link and len(self.feed.feed_link) > 6:
            url = urlparse.urljoin(self.feed.feed_link, 'favicon.ico')
        if not url: return None, None, None

        image, image_file = self.get_image_from_url(url)
        if not image:
            url = urlparse.urljoin(self.feed.feed_link, '/favicon.ico')
            image, image_file = self.get_image_from_url(url)
        # print 'Found: %s - %s' % (url, image)
        return image, image_file, url

    def get_image_from_url(self, url):
        # print 'Requesting: %s' % url
        if not url:
            return None, None

        @timelimit(30)
        def _1(url):
            headers = {
                'User-Agent': 'NewsBlur Favicon Fetcher - %s subscriber%s - %s '
                              '(Mozilla/5.0 (Macintosh; Intel Mac OS X 10_7_1) '
                              'AppleWebKit/534.48.3 (KHTML, like Gecko) Version/5.1 '
                              'Safari/534.48.3)' % (
                    self.feed.num_subscribers,
                    's' if self.feed.num_subscribers != 1 else '',
                    self.feed.permalink,
                ),
                'Connection': 'close',
            }
            try:
                request = urllib2.Request(url, headers=headers)
                icon = urllib2.urlopen(request).read()
            except Exception:
                return None
            return icon
        try:
            icon = _1(url)
        except TimeoutError:
            return None, None

        try:
            icon_file = StringIO(icon)
            image = Image.open(icon_file)
        except (IOError, ValueError):
            return None, None

        return image, icon_file

    def _url_from_html(self, content):
        url = None
        if not content: return url
        try:
            if isinstance(content, unicode):
                content = content.encode('utf-8')
            icon_path = lxml.html.fromstring(content).xpath(
                '//link[@rel="icon" or @rel="shortcut icon"]/@href'
            )
        except (lxml.etree.ParserError, TypeError):
            return url

        if icon_path:
            if str(icon_path[0]).startswith('http'):
                url = icon_path[0]
            else:
                url = urlparse.urljoin(self.feed.feed_link, icon_path[0])
        return url

    def normalize_image(self, image):
        # if image.size != (16, 16):
        #     image = image.resize((16, 16), Image.BICUBIC)
        if image.mode != 'RGBA':
            try:
                image = image.convert('RGBA')
            except IOError:
                pass

        return image

    def determine_dominant_color_in_image(self, image):
        NUM_CLUSTERS = 5

        # Convert image into array of values for each point.
        if image.mode == '1':
            image.convert('L')
        ar = numpy.array(image)
        # ar = scipy.misc.fromimage(image)
        shape = ar.shape

        # Reshape array of values to merge color bands. [[R], [G], [B], [A]] => [R, G, B, A]
        if len(shape) > 2:
            ar = ar.reshape(scipy.product(shape[:2]), shape[2])

        # Get NUM_CLUSTERS worth of centroids.
        codes, _ = scipy.cluster.vq.kmeans(ar, NUM_CLUSTERS)

        # Pare centroids, removing blacks and whites and shades of really dark and really light.
        original_codes = codes
        for low, hi in [(60, 200), (35, 230), (10, 250)]:
            codes = scipy.array([code for code in codes
                                 if not ((code[0] < low and code[1] < low and code[2] < low) or
                                         (code[0] > hi and code[1] > hi and code[2] > hi))])
            if not len(codes): codes = original_codes
            else: break

        # Assign codes (vector quantization). Each vector is compared to the centroids
        # and assigned the nearest one.
        vecs, _ = scipy.cluster.vq.vq(ar, codes)

        # Count occurences of each clustered vector.
        counts, bins = scipy.histogram(vecs, len(codes))

        # Show colors for each code in its hex value.
        # colors = [''.join(chr(c) for c in code).encode('hex') for code in codes]
        # total = scipy.sum(counts)
        # print dict(zip(colors, [count/float(total) for count in counts]))

        # Find the most frequent color, based on the counts.
        index_max = scipy.argmax(counts)
        peak = codes[index_max]
        color = ''.join(chr(c) for c in peak).encode('hex')

        return color[:6]

    def string_from_image(self, image):
        output = StringIO()
        image.save(output, 'png', quality=95)
        contents = output.getvalue()
        output.close()
        return contents.encode('base64')
    <|endoftext|>