<commit_before>import os
import hmac
import json
import time
import asyncio
import hashlib
import functools

import furl

from waterbutler.core import streams
from waterbutler.core import provider
from waterbutler.core import exceptions

from waterbutler.cloudfiles import settings
from waterbutler.cloudfiles.metadata import CloudFilesFileMetadata
from waterbutler.cloudfiles.metadata import CloudFilesFolderMetadata

def ensure_connection(func):
    """Runs ``_ensure_connection`` before continuing to the method
    """
    @functools.wraps(func)
    @asyncio.coroutine
    def wrapped(self, *args, **kwargs):
        yield from self._ensure_connection()
        return (yield from func(self, *args, **kwargs))
    return wrapped

class CloudFilesProvider(provider.BaseProvider):
    """Provider for Rackspace CloudFiles
    """

    def __init__(self, auth, credentials, settings):
        super().__init__(auth, credentials, settings)
        self.token = None
        self.endpoint = None
        self.temp_url_key = None
        self.region = self.credentials['region']
        self.og_token = self.credentials['token']
        self.username = self.credentials['username']
        self.container = self.settings['container']

    def can_intra_copy(self, dest_provider):
        return self == dest_provider

    def can_intra_move(self, dest_provider):
        return self == dest_provider

    @asyncio.coroutine
    def intra_copy(self, dest_provider, source_options, dest_options):
        url = dest_provider.build_url(dest_options['path'])
        yield from self.make_request(
            'PUT',
            url,
            headers={
                'X-Copy-From': os.path.join(self.container, source_options['path'])
            },
            expects=(201, ),
            throws=exceptions.IntraCopyError,
        )
        return (yield from dest_provider.metadata(dest_options['path']))

    @asyncio.coroutine
    def get_token(self):
        """Fetchs an access token from cloudfiles for actual api requests
        Returns the entire json response from the tokens endpoint
        Notably containing our token and proper endpoint to send requests to
        :rtype dict:
        """
        resp = yield from self.make_request(
            'POST',
            settings.AUTH_URL,
            data=json.dumps({
                'auth': {
                    'RAX-KSKEY:apiKeyCredentials': {
                        'username': self.username,
                        'apiKey': self.og_token,
                    }
                }
            }),
            headers={
                'Content-Type': 'application/json',
            },
            expects=(200, ),
        )
        data = yield from resp.json()
        return data

    @asyncio.coroutine
    def _ensure_connection(self):
        """Defines token, endpoint and temp_url_key if they are not already defined
        :raises ProviderError: If no temp url key is available
        """
        # Must have a temp url key for download and upload
        # Currently You must have one for everything however
        if not self.token or not self.endpoint or not self.temp_url_key:
            data = yield from self.get_token()
            self.token = data['access']['token']['id']
            self.endpoint = self.extract_endpoint(data)
            resp = yield from self.make_request('HEAD', self.endpoint, expects=(204, ))
            try:
                self.temp_url_key = resp.headers['X-Account-Meta-Temp-URL-Key'].encode()
            except KeyError:
                raise exceptions.ProviderError('No temp url key is available', code=503)

    @property
    def default_headers(self):
        return {
            'X-Auth-Token': self.token,
            'Accept': 'application/json',
        }

    def extract_endpoint(self, data):
        """Pulls the proper cloudfiles url from the return of tokens
        Very optimized.
        :param dict data: The json response from the token endpoint
        :rtype str:
        """
        for service in reversed(data['access']['serviceCatalog']):
            if service['name'].lower() == 'cloudfiles':
                for region in service['endpoints']:
                    if region['region'].lower() == self.region.lower():
                        return region['publicURL']

    def build_url(self, stream):
        """Build the url for the specified object
        :param str stream: The stream in question
        :rtype str:
        """
        url = furl.furl(self.endpoint)
        url.path.add(self.container)
        url.path.add(stream)
        return url.url

    def generate_url(self, stream, method='GET', seconds=settings.TEMP_URL_SECS):
        """Build and sign a temp url for the specified stream
        :param str stream: The requested stream's path
        :param str method: The HTTP method used to access the returned url
        :param int seconds: Time for the url to live
        :rtype str:
        """
        method = method.upper()
        expires = str(int(time.time() + seconds))
        url = furl.furl(self.build_url(stream))

        body = '\n'.join([method, expires, str(url.path)]).encode()
        signature = hmac.new(self.temp_url_key, body, hashlib.sha1).hexdigest()

        url.args.update({
            'temp_url_sig': signature,
            'temp_url_expires': expires,
        })
        return url.url

    @ensure_connection
    def download(self, path, accept_url=False, **kwargs):
        """Returns a ResponseStreamReader (Stream) for the specified path
        :param str path: Path to the object you want to download
        :param dict **kwargs: Additional arguments that are ignored
        :rtype str:
        :rtype ResponseStreamReader:
        :raises: waterbutler.FileNotFoundError
        """
        url = self.generate_url(path)

        if accept_url:
            return url

        resp = yield from self.make_request(
            'GET',
            url,
            expects=(200, ),
            throws=exceptions.DownloadError,
        )
        return streams.ResponseStreamReader(resp)

    @ensure_connection
    def upload(self, stream, path, **kwargs):
        """Uploads the given stream to S3
        :param ResponseStreamReader stream: The stream to put to Cloudfiles
        :param str path: The full path of the object to upload to/into
        :rtype ResponseStreamReader:
        """
        url = self.generate_url(path, 'PUT')
        yield from self.make_request(
            'PUT',
            url,
            data=stream,
            headers={'Content-Length': str(stream.size)},
            expects=(200, 201),
            throws=exceptions.UploadError,
        )
        return (yield from self.metadata(path))

    @ensure_connection
    def delete(self, path, **kwargs):
        """Deletes the key at the specified path
        :param str path: The path of the key to delete
        :rtype ResponseStreamReader:
        """
        yield from self.make_request(
            'DELETE',
            self.build_url(path),
            expects=(204, ),
            throws=exceptions.DeleteError,
        )

    @ensure_connection
    def metadata(self, path, **kwargs):
        """Get Metadata about the requested file or folder
        :param str path: The path to a key or folder
        :rtype dict:
        :rtype list:
        """
        url = furl.furl(self.build_url(''))
        url.args.update({'prefix': path, 'delimiter': '/'})
        resp = yield from self.make_request(
            'GET',
            url.url,
            expects=(200, 204),
            throws=exceptions.MetadataError,
        )

        data = yield from resp.json()

        if not data:
            raise exceptions.MetadataError(
                'Could not retrieve file or directory {0}'.format(path),
                code=404,
            )

        if path.endswith('/'):
            return [self._serialize_metadata(item) for item in data]
        if data:
            return self._serialize_metadata(data[0])

    def _serialize_metadata(self, metadata):
        if metadata.get('subdir'):
            return CloudFilesFolderMetadata(metadata).serialized()
        return CloudFilesFileMetadata(metadata).serialized()
<commit_msg>std metadata<commit_after>import os
import hmac
import json
import time
import asyncio
import hashlib
import functools

import furl

from waterbutler.core import streams
from waterbutler.core import provider
from waterbutler.core import exceptions

from waterbutler.cloudfiles import metadata
from waterbutler.cloudfiles import settings

def ensure_connection(func):
    """Runs ``_ensure_connection`` before continuing to the method
    """
    @functools.wraps(func)
    @asyncio.coroutine
    def wrapped(self, *args, **kwargs):
        yield from self._ensure_connection()
        return (yield from func(self, *args, **kwargs))
    return wrapped

class CloudFilesProvider(provider.BaseProvider):
    """Provider for Rackspace CloudFiles
    """

    def __init__(self, auth, credentials, settings):
        super().__init__(auth, credentials, settings)
        self.token = None
        self.endpoint = None
        self.temp_url_key = None
        self.region = self.credentials['region']
        self.og_token = self.credentials['token']
        self.username = self.credentials['username']
        self.container = self.settings['container']

    def can_intra_copy(self, dest_provider):
        return self == dest_provider

    def can_intra_move(self, dest_provider):
        return self == dest_provider

    @asyncio.coroutine
    def intra_copy(self, dest_provider, source_options, dest_options):
        url = dest_provider.build_url(dest_options['path'])
        yield from self.make_request(
            'PUT',
            url,
            headers={
                'X-Copy-From': os.path.join(self.container, source_options['path'])
            },
            expects=(201, ),
            throws=exceptions.IntraCopyError,
        )
        return (yield from dest_provider.metadata(dest_options['path']))

    @asyncio.coroutine
    def get_token(self):
        """Fetchs an access token from cloudfiles for actual api requests
        Returns the entire json response from the tokens endpoint
        Notably containing our token and proper endpoint to send requests to
        :rtype dict:
        """
        resp = yield from self.make_request(
            'POST',
            settings.AUTH_URL,
            data=json.dumps({
                'auth': {
                    'RAX-KSKEY:apiKeyCredentials': {
                        'username': self.username,
                        'apiKey': self.og_token,
                    }
                }
            }),
            headers={
                'Content-Type': 'application/json',
            },
            expects=(200, ),
        )
        data = yield from resp.json()
        return data

    @asyncio.coroutine
    def _ensure_connection(self):
        """Defines token, endpoint and temp_url_key if they are not already defined
        :raises ProviderError: If no temp url key is available
        """
        # Must have a temp url key for download and upload
        # Currently You must have one for everything however
        if not self.token or not self.endpoint or not self.temp_url_key:
            data = yield from self.get_token()
            self.token = data['access']['token']['id']
            self.endpoint = self.extract_endpoint(data)
            resp = yield from self.make_request('HEAD', self.endpoint, expects=(204, ))
            try:
                self.temp_url_key = resp.headers['X-Account-Meta-Temp-URL-Key'].encode()
            except KeyError:
                raise exceptions.ProviderError('No temp url key is available', code=503)

    @property
    def default_headers(self):
        return {
            'X-Auth-Token': self.token,
            'Accept': 'application/json',
        }

    def extract_endpoint(self, data):
        """Pulls the proper cloudfiles url from the return of tokens
        Very optimized.
        :param dict data: The json response from the token endpoint
        :rtype str:
        """
        for service in reversed(data['access']['serviceCatalog']):
            if service['name'].lower() == 'cloudfiles':
                for region in service['endpoints']:
                    if region['region'].lower() == self.region.lower():
                        return region['publicURL']

    def build_url(self, stream):
        """Build the url for the specified object
        :param str stream: The stream in question
        :rtype str:
        """
        url = furl.furl(self.endpoint)
        url.path.add(self.container)
        url.path.add(stream)
        return url.url

    def generate_url(self, stream, method='GET', seconds=settings.TEMP_URL_SECS):
        """Build and sign a temp url for the specified stream
        :param str stream: The requested stream's path
        :param str method: The HTTP method used to access the returned url
        :param int seconds: Time for the url to live
        :rtype str:
        """
        method = method.upper()
        expires = str(int(time.time() + seconds))
        url = furl.furl(self.build_url(stream))

        body = '\n'.join([method, expires, str(url.path)]).encode()
        signature = hmac.new(self.temp_url_key, body, hashlib.sha1).hexdigest()

        url.args.update({
            'temp_url_sig': signature,
            'temp_url_expires': expires,
        })
        return url.url

    @ensure_connection
    def download(self, path, accept_url=False, **kwargs):
        """Returns a ResponseStreamReader (Stream) for the specified path
        :param str path: Path to the object you want to download
        :param dict **kwargs: Additional arguments that are ignored
        :rtype str:
        :rtype ResponseStreamReader:
        :raises: waterbutler.FileNotFoundError
        """
        url = self.generate_url(path)

        if accept_url:
            return url

        resp = yield from self.make_request(
            'GET',
            url,
            expects=(200, ),
            throws=exceptions.DownloadError,
        )
        return streams.ResponseStreamReader(resp)

    @ensure_connection
    def upload(self, stream, path, **kwargs):
        """Uploads the given stream to S3
        :param ResponseStreamReader stream: The stream to put to Cloudfiles
        :param str path: The full path of the object to upload to/into
        :rtype ResponseStreamReader:
        """
        url = self.generate_url(path, 'PUT')
        yield from self.make_request(
            'PUT',
            url,
            data=stream,
            headers={'Content-Length': str(stream.size)},
            expects=(200, 201),
            throws=exceptions.UploadError,
        )
        return (yield from self.metadata(path))

    @ensure_connection
    def delete(self, path, **kwargs):
        """Deletes the key at the specified path
        :param str path: The path of the key to delete
        :rtype ResponseStreamReader:
        """
        yield from self.make_request(
            'DELETE',
            self.build_url(path),
            expects=(204, ),
            throws=exceptions.DeleteError,
        )

    @ensure_connection
    def metadata(self, path, **kwargs):
        """Get Metadata about the requested file or folder
        :param str path: The path to a key or folder
        :rtype dict:
        :rtype list:
        """
        url = furl.furl(self.build_url(''))
        url.args.update({'prefix': path, 'delimiter': '/'})
        resp = yield from self.make_request(
            'GET',
            url.url,
            expects=(200, 204),
            throws=exceptions.MetadataError,
        )

        data = yield from resp.json()

        if not data:
            raise exceptions.MetadataError(
                'Could not retrieve file or directory {0}'.format(path),
                code=404,
            )

        if path.endswith('/'):
            return [self._serialize_metadata(item) for item in data]
        if data:
            return self._serialize_metadata(data[0])

    def _serialize_metadata(self, data):
        if data.get('subdir'):
            return metadata.CloudFilesFolderMetadata(data).serialized()
        return metadata.CloudFilesFileMetadata(data).serialized()
<|endoftext|>


<commit_before>__version__ = (2, 0, 11)
<commit_msg>vbump<commit_after>__version__ = (2, 0, 12)
<|endoftext|>


<commit_before>from lxml import etree
import xmlschema
import types

class UnknownXType(Exception):

    pass

class XType(object):

    def __init__(self, pythonType, forceList = False):
        self.pythonType = pythonType
        self.forceList = forceList

def XTypeFromXObjectType(xObjectType):

    if (type(xObjectType) == type and
            issubclass(xObjectType, XObject)):
        return XType(xObjectType)
    elif issubclass(xObjectType.__class__, XType):
        return XType
    elif xObjectType == int:
        return XType(XObjectInt)
    elif xObjectType == str:
        return XType(str)
    elif type(xObjectType) == list:
        assert(len(xObjectType) == 1)
        return XType(XTypeFromXObjectType(xObjectType[0]).pythonType,
                     forceList = True)

    raise UnknownXType

class XObject(object):

    def _setAttribute(self, key, val):
        expectedType = getattr(self.__class__, key, None)
        if expectedType:
            expectedXType = XTypeFromXObjectType(expectedType)
            val = expectedXType.pythonType(val)
        else:
            expectedXType = None
            val = XObjectStr(val)

        self._setItem(key, val, expectedXType)

    def _setItem(self, key, val, xType = None):
        current = getattr(self, key, None)
        if xType and xType.forceList:
            # force the item to be a list, and use the type inside of
            # this list as the type of elements of the list
            if id(current) == id(xType):
                current = []
                setattr(self, key, current)

        if key not in self.__dict__:
            # This has not yet been set in the instance.
            setattr(self, key, val)
        elif type(current) == list:
            current.append(val)
        else:
            setattr(self, key, [ current, val ])

    def __init__(self, text):
        self.text = text

class XObjectInt(XObject, int):

    pass

class XObjectStr(XObject, str):

    pass

class XObjParseException(Exception):

    pass

def smiter(item):
    if hasattr(item, '__iter__'):
        return item
    else:
        return [ item ]

def parse(xml, rootXClass = None, nameSpaceMap = {}):

    def nsmap(s):
        for key, val in nameSpaceMap.iteritems():
            if s.startswith(key):
                s = s.replace(key, val)

        return s

    def parseElement(element, xClass = None):
        # handle the text for this tag
        if xClass:
            xobj = xClass(element.text)
        else:
            localTag = nsmap(element.tag)
            # create a subclass for this type
            NewClass = type(localTag + '_XObj_Type', (XObjectStr,), {})
            xobj = NewClass(element.text)

        # handle children
        for childElement in element.getchildren():
            if types.BuiltinFunctionType == type(childElement.tag):
                # this catches comments. this is ugly.
                continue

            tag = nsmap(childElement.tag)
            childType = getattr(xClass, tag, None)
            if childType:
                #import epdb;epdb.st()
                childXType = XTypeFromXObjectType(childType)
            else:
                childXType = None

            child = parseElement(childElement, xClass = childType)
            xobj._setItem(tag, child, childXType)

        # handle attributes
        for (key, val) in element.items():
            key = nsmap(key)
            xobj._setAttribute(key, val)

        # anything which is the same as in the class wasn't set in XML, so
        # set it to None
        for key, val in xobj.__class__.__dict__.items():
            if key[0] == '_': continue
            if getattr(xobj, key) == val:
                setattr(xobj, key, None)

        return xobj

    return parseElement(xml.getroot(), xClass = rootXClass)

def parsef(f, rootXClass = None, nameSpaceMap = {}):
    schemaf = None
    if schemaf:
        schemaXml = etree.parse(schemaf)
        schemaXObj = parse(schemaXml, nameSpaceMap =
                           { '{http://www.w3.org/2001/XMLSchema}' : 'xsd_' })
        schema = xmlschema.Schema(schemaXObj)
        schemaObj = etree.XMLSchema(schemaXml)
    else:
        schema = schemaXml = schemaXObj = schemaObj = None

    parser = etree.XMLParser(schema = schemaObj)
    xml = etree.parse(f, parser)

    return parse(xml, rootXClass = rootXClass, nameSpaceMap = nameSpaceMap)
<commit_msg>parseElement() takes an XType, not a python type<commit_after>from lxml import etree
import xmlschema
import types

class UnknownXType(Exception):

    pass

class XType(object):

    def __init__(self, pythonType, forceList = False):
        self.pythonType = pythonType
        self.forceList = forceList

def XTypeFromXObjectType(xObjectType):

    if (type(xObjectType) == type and
            issubclass(xObjectType, XObject)):
        return XType(xObjectType)
    elif issubclass(xObjectType.__class__, XType):
        return XType
    elif xObjectType == int:
        return XType(XObjectInt)
    elif xObjectType == str:
        return XType(str)
    elif type(xObjectType) == list:
        assert(len(xObjectType) == 1)
        return XType(XTypeFromXObjectType(xObjectType[0]).pythonType,
                     forceList = True)

    raise UnknownXType

class XObject(object):

    def _setAttribute(self, key, val):
        expectedType = getattr(self.__class__, key, None)
        if expectedType:
            expectedXType = XTypeFromXObjectType(expectedType)
            val = expectedXType.pythonType(val)
        else:
            expectedXType = None
            val = XObjectStr(val)

        self._setItem(key, val, expectedXType)

    def _setItem(self, key, val, xType = None):
        current = getattr(self, key, None)
        if xType and xType.forceList:
            # force the item to be a list, and use the type inside of
            # this list as the type of elements of the list
            if key not in self.__dict__:
                current = []
                setattr(self, key, current)

        if key not in self.__dict__:
            # This has not yet been set in the instance.
            setattr(self, key, val)
        elif type(current) == list:
            current.append(val)
        else:
            setattr(self, key, [ current, val ])

    def __init__(self, text):
        self.text = text

class XObjectInt(XObject, int):

    pass

class XObjectStr(XObject, str):

    pass

class XObjParseException(Exception):

    pass

def smiter(item):
    if hasattr(item, '__iter__'):
        return item
    else:
        return [ item ]

def parse(xml, rootXClass = None, nameSpaceMap = {}):

    def nsmap(s):
        for key, val in nameSpaceMap.iteritems():
            if s.startswith(key):
                s = s.replace(key, val)

        return s

    def parseElement(element, xType = None):
        # handle the text for this tag
        if xType:
            xobj = xType.pythonType(element.text)
        else:
            localTag = nsmap(element.tag)
            # create a subclass for this type
            NewClass = type(localTag + '_XObj_Type', (XObjectStr,), {})
            xobj = NewClass(element.text)

        # handle children
        for childElement in element.getchildren():
            if types.BuiltinFunctionType == type(childElement.tag):
                # this catches comments. this is ugly.
                continue

            tag = nsmap(childElement.tag)

            childXType = None
            if xType:
                childType = getattr(xType.pythonType, tag, None)
                if childType:
                    childXType = XTypeFromXObjectType(childType)

            child = parseElement(childElement, xType = childXType)
            xobj._setItem(tag, child, childXType)

        # handle attributes
        for (key, val) in element.items():
            key = nsmap(key)
            xobj._setAttribute(key, val)

        # anything which is the same as in the class wasn't set in XML, so
        # set it to None
        for key, val in xobj.__class__.__dict__.items():
            if key[0] == '_': continue
            if getattr(xobj, key) == val:
                setattr(xobj, key, None)

        return xobj

    if rootXClass:
        rootXType = XTypeFromXObjectType(rootXClass)
    else:
        rootXType = None

    return parseElement(xml.getroot(), xType = rootXType)

def parsef(f, rootXClass = None, nameSpaceMap = {}):
    schemaf = None
    if schemaf:
        schemaXml = etree.parse(schemaf)
        schemaXObj = parse(schemaXml, nameSpaceMap =
                           { '{http://www.w3.org/2001/XMLSchema}' : 'xsd_' })
        schema = xmlschema.Schema(schemaXObj)
        schemaObj = etree.XMLSchema(schemaXml)
    else:
        schema = schemaXml = schemaXObj = schemaObj = None

    parser = etree.XMLParser(schema = schemaObj)
    xml = etree.parse(f, parser)

    return parse(xml, rootXClass = rootXClass, nameSpaceMap = nameSpaceMap)
<|endoftext|>


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

# Copyright (c) 2015, Kaoru Esashika
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.

# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.

# * Neither the name of yacoinado nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import os
import re
import hashlib
import argparse
import requests
import bencodepy

COINADO_SECRET = os.getenv('YACOINADO_SECRET')

def get_torrent_file(url):
	response = requests.get(url)
	return response.content

def get_hash(binary):
	decoded_data = bencodepy.decode(binary)
	decoded_info = decoded_data[b'info']
	encoded_info = bencodepy.encode(decoded_info)
	torrent_hash = hashlib.sha1(encoded_info).hexdigest()
	return torrent_hash

def get_endpoint(torrent_hash, select=None):
	if select:
		endpoint = 'https://coinado.io/i/{}/search/{}?u={}'.format(torrent_hash, select, COINADO_SECRET)
	else:
		endpoint = 'https://coinado.io/i/{}/auto?u={}'.format(torrent_hash, COINADO_SECRET)
	return endpoint

def get_payload(torrent_hash, dest_path, select=None, permit_overwrite=False):

	def spinning_cursor():
		while True:
			for cursor in r'|/-\\':
				yield cursor
	spinner = spinning_cursor()

	def display_progress(count, size, total):
		percentage = count * size / total
		sys.stdout.write(
			'\r{} {:.3%} ({:,}/{:,} bytes) downloaded.'
			.format(next(spinner), percentage, count * size, total))
		sys.stdout.flush()

	endpoint = get_endpoint(torrent_hash, select)
	print('Connecting...')
	response = requests.get(endpoint, stream=True)
	if os.path.isdir(dest_path):
		match = re.search(r'''filename=['"](.*?)['"]''', response.headers['Content-Disposition'])
		filename = match.group(1)
		path = os.path.join(os.path.abspath(dest_path), filename)
	elif os.path.isfile(dest_path):
		if permit_overwrite:
			path = os.path.abspath(dest_path)
		else:
			print('file: {} exists.'.format(dest_path), file=sys.stderr)
			print('To overwrite the file, please add --force option', file=sys.stderr)
			return
	else:
		path = os.path.abspath(dest_path)

	with open(path, mode='wb') as fd:
		print('Connected.')
		print('Destination: {}'.format(path))
		for block in response.iter_content(1024):
			fd.write(block)
			display_progress(fd.tell(), 1, int(response.headers['Content-Length']))
	print('\nDownload completed.')

def get_source_list(args):
	if args is None:
		return list()
	lst = list()
	if '-' in args:
		for url in sys.stdin.readlines():
			lst.append(url)
	else:
		lst.extend(args)
	return lst

def url_to_hash(url):
	rawdata = get_torrent_file(url)
	return get_hash(rawdata)

def magnet_to_hash(magnet):
	match = re.search(r'xt=urn:btih:(.*?)(?:&.*)?$', magnet)
	return match.group(1)

def file_to_hash(path):
	with open(path, mode='rb') as fd:
		rawdata = fd.read()
	return get_hash(rawdata)

def source_to_hash(source):
	if source.startswith('http'):
		torrent_hash = url_to_hash(source)
	elif 'xt=urn:btih:' in source:
		torrent_hash = magnet_to_hash(source)
	elif os.path.exists(source):
		torrent_hash = file_to_hash(source)
	elif re.match('^[0-9a-fA-F]+$', source):
		torrent_hash = source
	else:
		print('Unknown source detected.', file=sys.stderr)
	return torrent_hash

def show_balance():
	response = requests.get('https://coinado.io/balance?u={}'.format(COINADO_SECRET))
	address = re.search(r'Bitcoin address: (?P<address>1[0-9a-zA-Z]+)', response.text, re.MULTILINE)
	balance = re.search(r'Balance: (?P<balance>.*)$', response.text, re.MULTILINE)
	pricing = re.search(r'Pricing: (?P<pricing>.*)$', response.text, re.MULTILINE)
	qrcode_image = requests.get(
		'https://chart.googleapis.com/chart?chs=250x250&cht=qr&chl={}'
		.format(address.group(1)))
	try:
		import sixel
	except ImportError:
		pass
	else:
		import io
		sixel_writer = sixel.SixelWriter()
		sixel_writer.draw(io.BytesIO(qrcode_image.content))
	print('Address: bitcoin:{}'.format(address.group(1)))
	print('QRcode: https://chart.googleapis.com/chart?chs=250x250&cht=qr&chl={}'.format(address.group(1)))
	print('Balance: {}'.format(balance.group(1)))
	print('Pricing: {}'.format(pricing.group(1)))

def dispatch_hash(args, t_hash):
	if args.filename:
		response = requests.head(get_endpoint(t_hash, args.select))
		match = re.search(r'''filename=['"](.*?)['"]''', response.headers['Content-Disposition'])
		print(match.group(1))
	elif args.endpoint:
		print(get_endpoint(t_hash, args.select))
	elif args.infohash:
		print(t_hash)
	else:
		if os.path.isfile(args.destination) and not args.force:
			print('To overwrite the file, please add --force option.', file=sys.stderr)
			return
		else:
			if args.quiet:
				sys.stdout = open(os.devnull, 'w')
			get_payload(t_hash, args.destination, args.select, args.force)

def setup_args():
	argparser = argparse.ArgumentParser()
	if len(sys.argv) == 1:
		argparser.print_usage()
	argparser.add_argument('torrent', nargs='*', help='url,hash,file,magnet...')
	argparser.add_argument('--quiet', action='store_true')
	g_source = argparser.add_argument_group(title='input option')
	g_source.add_argument('--select', help='specify the filename which you want')
	g_source.add_argument('--stdin', action='store_true')
	g_dest = argparser.add_argument_group(title='output option')
	g_dest.add_argument('--destination', default=os.getcwd(), help='save files to this directory')
	g_dest.add_argument('--force', action='store_true', help='permit overwrite')
	g_info = argparser.add_argument_group(title='infomation option')
	g_info = g_info.add_mutually_exclusive_group()
	g_info.add_argument('--inquiry', '--balance', action='store_true', help='show you account balance')
	g_info.add_argument('--filename', action='store_true', help='fetch filename')
	g_info.add_argument('--endpoint', action='store_true', help='fetch download url')
	g_info.add_argument('--infohash', action='store_true', help='calculate hash')
	args = argparser.parse_args()
	if args.torrent is None and args.stdin is not True:
		argparser.print_usage()
	return args

if __name__ == '__main__':
	args = setup_args()

	if args.inquiry:
		show_balance()
		sys.exit()

	for source in args.torrent:
		t_hash = source_to_hash(source)
		dispatch_hash(args, t_hash)

	if args.stdin:
		while True:
			try:
				source = input()
			except EOFError:
				sys.exit()
			t_hash = source_to_hash(source)
			dispatch_hash(args, t_hash)
<commit_msg>deal if coinado secret was not set<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*- #

# Copyright (c) 2015, Kaoru Esashika
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.

# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.

# * Neither the name of yacoinado nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import os
import re
import hashlib
import argparse
import requests
import bencodepy

COINADO_SECRET = os.getenv('YACOINADO_SECRET')
if not COINADO_SECRET:
	print('Please set YACOINADO_SECRET environment variable', file=sys.stderr)
	sys.exit()

def get_torrent_file(url):
	response = requests.get(url)
	return response.content

def get_hash(binary):
	decoded_data = bencodepy.decode(binary)
	decoded_info = decoded_data[b'info']
	encoded_info = bencodepy.encode(decoded_info)
	torrent_hash = hashlib.sha1(encoded_info).hexdigest()
	return torrent_hash

def get_endpoint(torrent_hash, select=None):
	if select:
		endpoint = 'https://coinado.io/i/{}/search/{}?u={}'.format(torrent_hash, select, COINADO_SECRET)
	else:
		endpoint = 'https://coinado.io/i/{}/auto?u={}'.format(torrent_hash, COINADO_SECRET)
	return endpoint

def get_payload(torrent_hash, dest_path, select=None, permit_overwrite=False):

	def spinning_cursor():
		while True:
			for cursor in r'|/-\\':
				yield cursor
	spinner = spinning_cursor()

	def display_progress(count, size, total):
		percentage = count * size / total
		sys.stdout.write(
			'\r{} {:.3%} ({:,}/{:,} bytes) downloaded.'
			.format(next(spinner), percentage, count * size, total))
		sys.stdout.flush()

	endpoint = get_endpoint(torrent_hash, select)
	print('Connecting...')
	response = requests.get(endpoint, stream=True)
	if os.path.isdir(dest_path):
		match = re.search(r'''filename=['"](.*?)['"]''', response.headers['Content-Disposition'])
		filename = match.group(1)
		path = os.path.join(os.path.abspath(dest_path), filename)
	elif os.path.isfile(dest_path):
		if permit_overwrite:
			path = os.path.abspath(dest_path)
		else:
			print('file: {} exists.'.format(dest_path), file=sys.stderr)
			print('To overwrite the file, please add --force option', file=sys.stderr)
			return
	else:
		path = os.path.abspath(dest_path)

	with open(path, mode='wb') as fd:
		print('Connected.')
		print('Destination: {}'.format(path))
		for block in response.iter_content(1024):
			fd.write(block)
			display_progress(fd.tell(), 1, int(response.headers['Content-Length']))
	print('\nDownload completed.')

def get_source_list(args):
	if args is None:
		return list()
	lst = list()
	if '-' in args:
		for url in sys.stdin.readlines():
			lst.append(url)
	else:
		lst.extend(args)
	return lst

def url_to_hash(url):
	rawdata = get_torrent_file(url)
	return get_hash(rawdata)

def magnet_to_hash(magnet):
	match = re.search(r'xt=urn:btih:(.*?)(?:&.*)?$', magnet)
	return match.group(1)

def file_to_hash(path):
	with open(path, mode='rb') as fd:
		rawdata = fd.read()
	return get_hash(rawdata)

def source_to_hash(source):
	if source.startswith('http'):
		torrent_hash = url_to_hash(source)
	elif 'xt=urn:btih:' in source:
		torrent_hash = magnet_to_hash(source)
	elif os.path.exists(source):
		torrent_hash = file_to_hash(source)
	elif re.match('^[0-9a-fA-F]+$', source):
		torrent_hash = source
	else:
		print('Unknown source detected.', file=sys.stderr)
	return torrent_hash

def show_balance():
	response = requests.get('https://coinado.io/balance?u={}'.format(COINADO_SECRET))
	address = re.search(r'Bitcoin address: (?P<address>1[0-9a-zA-Z]+)', response.text, re.MULTILINE)
	balance = re.search(r'Balance: (?P<balance>.*)$', response.text, re.MULTILINE)
	pricing = re.search(r'Pricing: (?P<pricing>.*)$', response.text, re.MULTILINE)
	qrcode_image = requests.get(
		'https://chart.googleapis.com/chart?chs=250x250&cht=qr&chl={}'
		.format(address.group(1)))
	try:
		import sixel
	except ImportError:
		pass
	else:
		import io
		sixel_writer = sixel.SixelWriter()
		sixel_writer.draw(io.BytesIO(qrcode_image.content))
	print('Address: bitcoin:{}'.format(address.group(1)))
	print('QRcode: https://chart.googleapis.com/chart?chs=250x250&cht=qr&chl={}'.format(address.group(1)))
	print('Balance: {}'.format(balance.group(1)))
	print('Pricing: {}'.format(pricing.group(1)))

def dispatch_hash(args, t_hash):
	if args.filename:
		response = requests.head(get_endpoint(t_hash, args.select))
		match = re.search(r'''filename=['"](.*?)['"]''', response.headers['Content-Disposition'])
		print(match.group(1))
	elif args.endpoint:
		print(get_endpoint(t_hash, args.select))
	elif args.infohash:
		print(t_hash)
	else:
		if os.path.isfile(args.destination) and not args.force:
			print('To overwrite the file, please add --force option.', file=sys.stderr)
			return
		else:
			if args.quiet:
				sys.stdout = open(os.devnull, 'w')
			get_payload(t_hash, args.destination, args.select, args.force)

def setup_args():
	argparser = argparse.ArgumentParser()
	if len(sys.argv) == 1:
		argparser.print_usage()
	argparser.add_argument('torrent', nargs='*', help='url,hash,file,magnet...')
	argparser.add_argument('--quiet', action='store_true')
	g_source = argparser.add_argument_group(title='input option')
	g_source.add_argument('--select', help='specify the filename which you want')
	g_source.add_argument('--stdin', action='store_true')
	g_dest = argparser.add_argument_group(title='output option')
	g_dest.add_argument('--destination', default=os.getcwd(), help='save files to this directory')
	g_dest.add_argument('--force', action='store_true', help='permit overwrite')
	g_info = argparser.add_argument_group(title='infomation option')
	g_info = g_info.add_mutually_exclusive_group()
	g_info.add_argument('--inquiry', '--balance', action='store_true', help='show you account balance')
	g_info.add_argument('--filename', action='store_true', help='fetch filename')
	g_info.add_argument('--endpoint', action='store_true', help='fetch download url')
	g_info.add_argument('--infohash', action='store_true', help='calculate hash')
	args = argparser.parse_args()
	if args.torrent is None and args.stdin is not True:
		argparser.print_usage()
	return args

if __name__ == '__main__':
	args = setup_args()

	if args.inquiry:
		show_balance()
		sys.exit()

	for source in args.torrent:
		t_hash = source_to_hash(source)
		dispatch_hash(args, t_hash)

	if args.stdin:
		while True:
			try:
				source = input()
			except EOFError:
				sys.exit()
			t_hash = source_to_hash(source)
			dispatch_hash(args, t_hash)
<|endoftext|>


<commit_before># Copyright (c) 2018 Yubico AB
# All rights reserved.
#
#   Redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following
#   conditions are met:
#
#    1. Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#    2. Redistributions in binary form must reproduce the above
#       copyright notice, this list of conditions and the following
#       disclaimer in the documentation and/or other materials provided
#       with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from __future__ import absolute_import

import time
import logging
from ctypes import sizeof, byref, c_uint, create_string_buffer
from .driver_otp import ykpers, check, YkpersError
from .util import time_challenge, parse_totp_hash, format_code, hmac_shorten_key
from .scancodes import encode, KEYBOARD_LAYOUT
from enum import IntEnum, unique
from binascii import a2b_hex, b2a_hex

logger = logging.getLogger(__name__)

@unique
class SLOT(IntEnum):
    CONFIG = 0x01
    CONFIG2 = 0x03
    UPDATE1 = 0x04
    UPDATE2 = 0x05
    SWAP = 0x06

SLOTS = [-1, 0x30, 0x38]
_ACCESS_CODE_LENGTH = 6
_RESET_ACCESS_CODE = b'\x00' * _ACCESS_CODE_LENGTH

def slot_to_cmd(slot, update=False):
    if slot == 1:
        return SLOT.UPDATE1 if update else SLOT.CONFIG
    elif slot == 2:
        return SLOT.UPDATE2 if update else SLOT.CONFIG2
    else:
        raise ValueError('slot must be 1 or 2')

class OtpController(object):

    def __init__(self, driver):
        self._driver = driver
        self._dev = driver.ykpers_dev
        self._access_code = None

    @property
    def access_code(self):
        return self._access_code

    @access_code.setter
    def access_code(self, value):
        self._access_code = value

    def _create_cfg(self, cmd):
        st = ykpers.ykds_alloc()
        cfg = ykpers.ykp_alloc()
        try:
            check(ykpers.yk_get_status(self._dev, st))
            ykpers.ykp_configure_version(cfg, st)
            ykpers.ykp_configure_command(cfg, cmd)
            check(ykpers.ykp_set_extflag(cfg, 'SERIAL_API_VISIBLE'))
            check(ykpers.ykp_set_extflag(cfg, 'ALLOW_UPDATE'))
            if self.access_code is not None:
                check(ykpers.ykp_set_access_code(
                    cfg, self.access_code, _ACCESS_CODE_LENGTH))
            return cfg
        except YkpersError:
            ykpers.ykp_free_config(cfg)
            raise
        finally:
            ykpers.ykds_free(st)

    @property
    def slot_status(self):
        return self._driver.slot_status

    def program_otp(self, slot, key, fixed, uid, append_cr=True):
        if len(key) != 16:
            raise ValueError('key must be 16 bytes')
        if len(uid) != 6:
            raise ValueError('private ID must be 6 bytes')
        if len(fixed) > 16:
            raise ValueError('public ID must be <= 16 bytes')

        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)

        try:
            check(ykpers.ykp_set_fixed(cfg, fixed, len(fixed)))
            check(ykpers.ykp_set_uid(cfg, uid, 6))
            ykpers.ykp_AES_key_from_raw(cfg, key)
            if append_cr:
                check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))
            check(ykpers.yk_write_command(self._dev,
                                          ykpers.ykp_core_config(cfg),
                                          cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def program_static(self, slot, password, append_cr=True,
                       keyboard_layout=KEYBOARD_LAYOUT.MODHEX):
        pw_len = len(password)
        if self._driver.version < (2, 0, 0):
            raise ValueError('static password requires YubiKey 2.0.0 or later')
        elif self._driver.version < (2, 2, 0) and pw_len > 16:
            raise ValueError('password too long, this device supports a '
                             'maximum of %d characters' % 16)
        elif pw_len > 38:
            raise ValueError('password too long, this device supports a '
                             'maximum of %d characters' % 38)

        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)

        try:
            check(ykpers.ykp_set_cfgflag(cfg, 'SHORT_TICKET'))

            if append_cr:
                check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))

            pw_bytes = encode(password, keyboard_layout=keyboard_layout)
            if pw_len <= 16:  # All in fixed
                check(ykpers.ykp_set_fixed(cfg, pw_bytes, pw_len))
            elif pw_len <= 16 + 6:  # All in fixed and uid
                check(ykpers.ykp_set_fixed(cfg, pw_bytes[:-6], pw_len - 6))
                check(ykpers.ykp_set_uid(cfg, pw_bytes[-6:], 6))
            else:  # All in fixed + uid + key
                check(ykpers.ykp_set_fixed(cfg, pw_bytes[:-22], pw_len - 22))
                check(ykpers.ykp_set_uid(cfg, pw_bytes[-22:-16], 6))
                ykpers.ykp_AES_key_from_raw(cfg, pw_bytes[-16:])

            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def program_chalresp(self, slot, key, touch=False):
        if self._driver.version < (2, 2, 0):
            raise ValueError('challenge-response requires YubiKey 2.2.0 or '
                             'later')
        key = hmac_shorten_key(key, 'SHA1')
        if len(key) > 20:
            raise ValueError('key lengths >20 bytes not supported')
        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)
        key = key.ljust(20, b'\0')  # Pad key to 20 bytes
        try:
            check(ykpers.ykp_set_tktflag(cfg, 'CHAL_RESP'))
            check(ykpers.ykp_set_cfgflag(cfg, 'CHAL_HMAC'))
            check(ykpers.ykp_set_cfgflag(cfg, 'HMAC_LT64'))
            if touch:
                check(ykpers.ykp_set_cfgflag(cfg, 'CHAL_BTN_TRIG'))
            ykpers.ykp_HMAC_key_from_raw(cfg, key)
            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def calculate(
            self, slot, challenge=None, totp=False,
            digits=6, wait_for_touch=True):
        if totp:
            if challenge is None:
                challenge = time_challenge(time.time())
            else:
                challenge = time_challenge(challenge)
        else:
            challenge = a2b_hex(challenge)
        # Always use 64 byte challenges, assume HMAC_LT64
        challenge = challenge.ljust(
            64, b'\1' if challenge.endswith(b'\0') else b'\0')
        resp = create_string_buffer(64)
        # Some versions of the NEO firmware returns error 11 too often.
        # Give the YubiKey 10 tries to do the calculation.
        for idx in range(10):
            try:
                logger.debug(
                    'Sending a challenge to the device. Slot %s. '
                    'Attempt %s. Wait for touch is %s.', slot, idx + 1,
                    wait_for_touch)
                check(ykpers.yk_challenge_response(
                        self._dev, SLOTS[slot], wait_for_touch,
                        len(challenge), challenge, sizeof(resp), resp))
            except YkpersError as e:
                if idx < 10 and e.errno == 11 and wait_for_touch is True:
                    # Error 11 when wait_for_touch is true is an unexpected
                    # state, let's try again.
                    continue
                elif wait_for_touch is False:
                    logger.debug('Got %s as expected.', e)
                    # NEOs and very old YK4s might still be blinking,
                    # lets try to read the serial to cancel it.
                    ykpers.yk_get_serial(self._dev, 0, 0, byref(c_uint()))
                    raise
                else:
                    logger.debug('YkpersError: %s', e)
                    raise
            # We got a result, break the loop.
            break
        if totp:
            return format_code(parse_totp_hash(resp.raw[:20]), digits)
        else:
            return b2a_hex(resp.raw[:20])

    def program_hotp(self, slot, key, imf=0, hotp8=False, append_cr=True):
        if self._driver.version < (2, 1, 0):
            raise ValueError('HOTP requires YubiKey 2.1.0 or later')
        key = hmac_shorten_key(key, 'SHA1')
        if len(key) > 20:
            raise ValueError('key lengths >20 bytes not supported')
        key = key.ljust(20, b'\0')  # Pad key to 20 bytes
        if imf % 16 != 0:
            raise ValueError('imf must be a multiple of 16')
        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)

        try:
            check(ykpers.ykp_set_tktflag(cfg, 'OATH_HOTP'))
            check(ykpers.ykp_set_oath_imf(cfg, imf))
            if hotp8:
                check(ykpers.ykp_set_cfgflag(cfg, 'OATH_HOTP8'))
            if append_cr:
                check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))
            ykpers.ykp_HMAC_key_from_raw(cfg, key)
            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def zap_slot(self, slot):
        check(ykpers.yk_write_command(self._dev, None, slot_to_cmd(slot),
                                      self.access_code))

    def swap_slots(self):
        if self._driver.version < (2, 3, 0):
            raise ValueError('swapping slots requires YubiKey 2.3.0 or later')
        cfg = self._create_cfg(SLOT.SWAP)
        try:
            ycfg = ykpers.ykp_core_config(cfg)
            check(ykpers.yk_write_command(self._dev, ycfg, SLOT.SWAP, None))
        finally:
            ykpers.ykp_free_config(cfg)

    def configure_ndef_slot(self, slot, prefix='https://my.yubico.com/yk/#'):
        ndef = ykpers.ykp_alloc_ndef()
        try:
            check(ykpers.ykp_construct_ndef_uri(ndef, prefix.encode()))
            check(ykpers.yk_write_ndef2(self._dev, ndef, slot))
        finally:
            ykpers.ykp_free_ndef(ndef)

    @property
    def _has_update_access_code_bug(self):
        return (4, 3, 1) < self._driver.version < (4, 3, 6)

    def set_access_code(self, slot, new_code=None, update=True,
                        allow_zero=False):
        if update and self._driver.version < (2, 3, 0):
            raise ValueError('Update requires YubiKey 2.3.0 or later')
        if not update and new_code is not None:
            raise ValueError('Cannot set new access code unless updating slot')
        if new_code == _RESET_ACCESS_CODE:
            raise ValueError('Cannot set access code to special value zero.')
        if new_code is not None and self._has_update_access_code_bug:
            raise ValueError(
                'This YubiKey firmware does not support updating the access '
                'code after programming the slot. Please set the access '
                'code when initially programming the slot instead.')

        cmd = slot_to_cmd(slot, update)
        cfg = self._create_cfg(cmd)
        try:
            check(ykpers.ykp_set_access_code(
                cfg, new_code or _RESET_ACCESS_CODE, _ACCESS_CODE_LENGTH))
            ycfg = ykpers.ykp_core_config(cfg)
            check(ykpers.yk_write_command(self._dev, ycfg, cmd,
                                          self.access_code))

            self.access_code = new_code

        finally:
            ykpers.ykp_free_config(cfg)

    def delete_access_code(self, slot):
        if self._has_update_access_code_bug:
            raise ValueError(
                'This YubiKey firmware does not support deleting the access '
                'code after programming the slot. Please delete and re-program '
                'the slot instead.')

        self.set_access_code(slot, None)

    def update_settings(self, slot, enter=True, pacing=None):
        cmd = slot_to_cmd(slot, update=True)
        cfg = self._create_cfg(cmd)
        if enter:
            check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))

        # Output speed throttling
        if pacing == 20:
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_10MS'))
        elif pacing == 40:
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_20MS'))
        elif pacing == 60:
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_10MS'))
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_20MS'))

        try:
            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    @property
    def is_in_fips_mode(self):
        return self._driver.is_in_fips_mode
<commit_msg>otp: update comment<commit_after># Copyright (c) 2018 Yubico AB
# All rights reserved.
#
#   Redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following
#   conditions are met:
#
#    1. Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#    2. Redistributions in binary form must reproduce the above
#       copyright notice, this list of conditions and the following
#       disclaimer in the documentation and/or other materials provided
#       with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from __future__ import absolute_import

import time
import logging
from ctypes import sizeof, byref, c_uint, create_string_buffer
from .driver_otp import ykpers, check, YkpersError
from .util import time_challenge, parse_totp_hash, format_code, hmac_shorten_key
from .scancodes import encode, KEYBOARD_LAYOUT
from enum import IntEnum, unique
from binascii import a2b_hex, b2a_hex

logger = logging.getLogger(__name__)

@unique
class SLOT(IntEnum):
    CONFIG = 0x01
    CONFIG2 = 0x03
    UPDATE1 = 0x04
    UPDATE2 = 0x05
    SWAP = 0x06

SLOTS = [-1, 0x30, 0x38]
_ACCESS_CODE_LENGTH = 6
_RESET_ACCESS_CODE = b'\x00' * _ACCESS_CODE_LENGTH

def slot_to_cmd(slot, update=False):
    if slot == 1:
        return SLOT.UPDATE1 if update else SLOT.CONFIG
    elif slot == 2:
        return SLOT.UPDATE2 if update else SLOT.CONFIG2
    else:
        raise ValueError('slot must be 1 or 2')

class OtpController(object):

    def __init__(self, driver):
        self._driver = driver
        self._dev = driver.ykpers_dev
        self._access_code = None

    @property
    def access_code(self):
        return self._access_code

    @access_code.setter
    def access_code(self, value):
        self._access_code = value

    def _create_cfg(self, cmd):
        st = ykpers.ykds_alloc()
        cfg = ykpers.ykp_alloc()
        try:
            check(ykpers.yk_get_status(self._dev, st))
            ykpers.ykp_configure_version(cfg, st)
            ykpers.ykp_configure_command(cfg, cmd)
            check(ykpers.ykp_set_extflag(cfg, 'SERIAL_API_VISIBLE'))
            check(ykpers.ykp_set_extflag(cfg, 'ALLOW_UPDATE'))
            if self.access_code is not None:
                check(ykpers.ykp_set_access_code(
                    cfg, self.access_code, _ACCESS_CODE_LENGTH))
            return cfg
        except YkpersError:
            ykpers.ykp_free_config(cfg)
            raise
        finally:
            ykpers.ykds_free(st)

    @property
    def slot_status(self):
        return self._driver.slot_status

    def program_otp(self, slot, key, fixed, uid, append_cr=True):
        if len(key) != 16:
            raise ValueError('key must be 16 bytes')
        if len(uid) != 6:
            raise ValueError('private ID must be 6 bytes')
        if len(fixed) > 16:
            raise ValueError('public ID must be <= 16 bytes')

        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)

        try:
            check(ykpers.ykp_set_fixed(cfg, fixed, len(fixed)))
            check(ykpers.ykp_set_uid(cfg, uid, 6))
            ykpers.ykp_AES_key_from_raw(cfg, key)
            if append_cr:
                check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))
            check(ykpers.yk_write_command(self._dev,
                                          ykpers.ykp_core_config(cfg),
                                          cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def program_static(self, slot, password, append_cr=True,
                       keyboard_layout=KEYBOARD_LAYOUT.MODHEX):
        pw_len = len(password)
        if self._driver.version < (2, 0, 0):
            raise ValueError('static password requires YubiKey 2.0.0 or later')
        elif self._driver.version < (2, 2, 0) and pw_len > 16:
            raise ValueError('password too long, this device supports a '
                             'maximum of %d characters' % 16)
        elif pw_len > 38:
            raise ValueError('password too long, this device supports a '
                             'maximum of %d characters' % 38)

        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)

        try:
            check(ykpers.ykp_set_cfgflag(cfg, 'SHORT_TICKET'))

            if append_cr:
                check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))

            pw_bytes = encode(password, keyboard_layout=keyboard_layout)
            if pw_len <= 16:  # All in fixed
                check(ykpers.ykp_set_fixed(cfg, pw_bytes, pw_len))
            elif pw_len <= 16 + 6:  # All in fixed and uid
                check(ykpers.ykp_set_fixed(cfg, pw_bytes[:-6], pw_len - 6))
                check(ykpers.ykp_set_uid(cfg, pw_bytes[-6:], 6))
            else:  # All in fixed + uid + key
                check(ykpers.ykp_set_fixed(cfg, pw_bytes[:-22], pw_len - 22))
                check(ykpers.ykp_set_uid(cfg, pw_bytes[-22:-16], 6))
                ykpers.ykp_AES_key_from_raw(cfg, pw_bytes[-16:])

            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def program_chalresp(self, slot, key, touch=False):
        if self._driver.version < (2, 2, 0):
            raise ValueError('challenge-response requires YubiKey 2.2.0 or '
                             'later')
        key = hmac_shorten_key(key, 'SHA1')
        if len(key) > 20:
            raise ValueError('key lengths >20 bytes not supported')
        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)
        key = key.ljust(20, b'\0')  # Pad key to 20 bytes
        try:
            check(ykpers.ykp_set_tktflag(cfg, 'CHAL_RESP'))
            check(ykpers.ykp_set_cfgflag(cfg, 'CHAL_HMAC'))
            check(ykpers.ykp_set_cfgflag(cfg, 'HMAC_LT64'))
            if touch:
                check(ykpers.ykp_set_cfgflag(cfg, 'CHAL_BTN_TRIG'))
            ykpers.ykp_HMAC_key_from_raw(cfg, key)
            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def calculate(
            self, slot, challenge=None, totp=False,
            digits=6, wait_for_touch=True):
        if totp:
            if challenge is None:
                challenge = time_challenge(time.time())
            else:
                challenge = time_challenge(challenge)
        else:
            challenge = a2b_hex(challenge)
        # Pad challenge when < 64 bytes
        challenge = challenge.ljust(
            64, b'\1' if challenge.endswith(b'\0') else b'\0')
        resp = create_string_buffer(64)
        # Some versions of the NEO firmware returns error 11 too often.
        # Give the YubiKey 10 tries to do the calculation.
        for idx in range(10):
            try:
                logger.debug(
                    'Sending a challenge to the device. Slot %s. '
                    'Attempt %s. Wait for touch is %s.', slot, idx + 1,
                    wait_for_touch)
                check(ykpers.yk_challenge_response(
                        self._dev, SLOTS[slot], wait_for_touch,
                        len(challenge), challenge, sizeof(resp), resp))
            except YkpersError as e:
                if idx < 10 and e.errno == 11 and wait_for_touch is True:
                    # Error 11 when wait_for_touch is true is an unexpected
                    # state, let's try again.
                    continue
                elif wait_for_touch is False:
                    logger.debug('Got %s as expected.', e)
                    # NEOs and very old YK4s might still be blinking,
                    # lets try to read the serial to cancel it.
                    ykpers.yk_get_serial(self._dev, 0, 0, byref(c_uint()))
                    raise
                else:
                    logger.debug('YkpersError: %s', e)
                    raise
            # We got a result, break the loop.
            break
        if totp:
            return format_code(parse_totp_hash(resp.raw[:20]), digits)
        else:
            return b2a_hex(resp.raw[:20])

    def program_hotp(self, slot, key, imf=0, hotp8=False, append_cr=True):
        if self._driver.version < (2, 1, 0):
            raise ValueError('HOTP requires YubiKey 2.1.0 or later')
        key = hmac_shorten_key(key, 'SHA1')
        if len(key) > 20:
            raise ValueError('key lengths >20 bytes not supported')
        key = key.ljust(20, b'\0')  # Pad key to 20 bytes
        if imf % 16 != 0:
            raise ValueError('imf must be a multiple of 16')
        cmd = slot_to_cmd(slot)
        cfg = self._create_cfg(cmd)

        try:
            check(ykpers.ykp_set_tktflag(cfg, 'OATH_HOTP'))
            check(ykpers.ykp_set_oath_imf(cfg, imf))
            if hotp8:
                check(ykpers.ykp_set_cfgflag(cfg, 'OATH_HOTP8'))
            if append_cr:
                check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))
            ykpers.ykp_HMAC_key_from_raw(cfg, key)
            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    def zap_slot(self, slot):
        check(ykpers.yk_write_command(self._dev, None, slot_to_cmd(slot),
                                      self.access_code))

    def swap_slots(self):
        if self._driver.version < (2, 3, 0):
            raise ValueError('swapping slots requires YubiKey 2.3.0 or later')
        cfg = self._create_cfg(SLOT.SWAP)
        try:
            ycfg = ykpers.ykp_core_config(cfg)
            check(ykpers.yk_write_command(self._dev, ycfg, SLOT.SWAP, None))
        finally:
            ykpers.ykp_free_config(cfg)

    def configure_ndef_slot(self, slot, prefix='https://my.yubico.com/yk/#'):
        ndef = ykpers.ykp_alloc_ndef()
        try:
            check(ykpers.ykp_construct_ndef_uri(ndef, prefix.encode()))
            check(ykpers.yk_write_ndef2(self._dev, ndef, slot))
        finally:
            ykpers.ykp_free_ndef(ndef)

    @property
    def _has_update_access_code_bug(self):
        return (4, 3, 1) < self._driver.version < (4, 3, 6)

    def set_access_code(self, slot, new_code=None, update=True,
                        allow_zero=False):
        if update and self._driver.version < (2, 3, 0):
            raise ValueError('Update requires YubiKey 2.3.0 or later')
        if not update and new_code is not None:
            raise ValueError('Cannot set new access code unless updating slot')
        if new_code == _RESET_ACCESS_CODE:
            raise ValueError('Cannot set access code to special value zero.')
        if new_code is not None and self._has_update_access_code_bug:
            raise ValueError(
                'This YubiKey firmware does not support updating the access '
                'code after programming the slot. Please set the access '
                'code when initially programming the slot instead.')

        cmd = slot_to_cmd(slot, update)
        cfg = self._create_cfg(cmd)
        try:
            check(ykpers.ykp_set_access_code(
                cfg, new_code or _RESET_ACCESS_CODE, _ACCESS_CODE_LENGTH))
            ycfg = ykpers.ykp_core_config(cfg)
            check(ykpers.yk_write_command(self._dev, ycfg, cmd,
                                          self.access_code))

            self.access_code = new_code

        finally:
            ykpers.ykp_free_config(cfg)

    def delete_access_code(self, slot):
        if self._has_update_access_code_bug:
            raise ValueError(
                'This YubiKey firmware does not support deleting the access '
                'code after programming the slot. Please delete and re-program '
                'the slot instead.')

        self.set_access_code(slot, None)

    def update_settings(self, slot, enter=True, pacing=None):
        cmd = slot_to_cmd(slot, update=True)
        cfg = self._create_cfg(cmd)
        if enter:
            check(ykpers.ykp_set_tktflag(cfg, 'APPEND_CR'))

        # Output speed throttling
        if pacing == 20:
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_10MS'))
        elif pacing == 40:
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_20MS'))
        elif pacing == 60:
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_10MS'))
            check(ykpers.ykp_set_cfgflag(cfg, 'PACING_20MS'))

        try:
            check(ykpers.yk_write_command(
                self._dev, ykpers.ykp_core_config(cfg), cmd, self.access_code))
        finally:
            ykpers.ykp_free_config(cfg)

    @property
    def is_in_fips_mode(self):
        return self._driver.is_in_fips_mode
<|endoftext|>


<commit_before># Copyright (c) 2011 Bogdan Popa
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
from pyglet import app, gl, text, window
from string import ascii_lowercase, digits

from ymage.helpers import reschedule, reschedule_once
from ymage.slideshow import Slideshow

class Printer(text.Label):
    def __init__(self):
        super(Printer, self).__init__(
            "", "Arial", 12, False, False,
            (0, 255, 0, 255), 10, 10
        )

    def _print(self, message, duration=5):
        self.text = message
        reschedule_once(self.clear, duration)

    def clear(self, dt=None):
        self.text = ""

class Reader(object):
    def __init__(self, printer):
        self.is_reading = False
        self.callback = None
        self.printer = printer
        self.message = ""
        self._input = ""

    def read(self, symbol, modifiers):
        self.is_reading = True
        if symbol == window.key.RETURN:
            self.printer.clear()
            self.toggle_reading()
            self.callback(self._input)
            return
        elif symbol == window.key.ESCAPE:
            self.printer.clear()
            self.toggle_reading()
            return
        elif symbol == window.key.BACKSPACE:
            self._input = self._input[:-1]
        try:
            self._input += {
                window.key.PERIOD: ".",
                window.key.MINUS: "-",
                window.key.SPACE: " ",
            }[symbol]
        except KeyError:
            pass
        representation = window.key.symbol_string(symbol)
        representation = representation.replace("_", "")
        if representation in (ascii_lowercase + digits):
            if modifiers & window.key.MOD_SHIFT:
                self._input += representation
            else:
                self._input += representation.lower()

        self.printer._print(
            "%s: %s_" % (self.message, self._input)
        )

    def start_reading(self, message, callback):
        self.callback = callback
        self.message = message
        self._input = ""
        self.printer._print(
            "%s: _" % self.message
        )
        self.toggle_reading()

    def toggle_reading(self):
        self.is_reading = not self.is_reading

class Window(window.Window):
    def __init__(self, options):
        super(Window, self).__init__(
            800, 600, "Loading...", True, None, False
        )
        self.slideshow = Slideshow(options)
        self.options = options
        self.printer = Printer()
        self.reader = Reader(self.printer)
        self.setup()

    def setup(self):
        if not self.options.windowed:
            self.toggle_fullscreen()

    def toggle_fullscreen(self):
        self.activate()
        self.set_mouse_visible(self.fullscreen)
        self.set_fullscreen(not self.fullscreen)

    def update_caption(self):
        self.set_caption("ymage [%d/%d]" % (
            self.slideshow.index + 1, len(self.slideshow.slides)
        ))

    # Actions
    def _set_duration(self):
        self.reader.start_reading(
            "Duration", self.slideshow.set_duration
        )

    def _increase_duration(self):
        self.printer._print("Duration: {}".format(
            self.slideshow.increase_duration()
        ))

    def _decrease_duration(self):
        self.printer._print("Duration: {}".format(
            self.slideshow.decrease_duration()
        ))

    def _previous_slide(self):
        self.slideshow.display(action="previous")

    def _next_slide(self):
        self.slideshow.display(action="next")

    def _toggle_paused(self):
        self.printer._print(
            self.slideshow.toggle_paused()
        )

    def _random(self):
        self.slideshow.display(action="random", previous=False)

    def _previous_random(self):
        self.slideshow.display(action="random", previous=True)

    def _print_info(self):
        self.printer._print("{} [{}/{}]".format(
            self.slideshow.get_current(),
            self.slideshow.index + 1,
            len(self.slideshow.slides)
        ))

    def _search(self, query=None):
        if query is None:
            self.reader.start_reading("Search for", self.search)
        else:
            self.slideshow.display(action="search", query=query)

    def _jump(self, index=None):
        if index is None:
            self.reader.start_reading("Slide", self.jump)
        else:
            self.slideshow.display(action="jump", index=index)

    def on_draw(self):
        self.clear()
        try:
            self.slideshow.draw_slide(self.width, self.height)
        except gl.lib.GLException:
            # In case one of the slides is corrupted
            # move on to the next
            self.slideshow.display()
        self.update_caption()
        self.printer.draw()

    def on_key_press(self, symbol, modifiers):
        if self.reader.is_reading:
            self.reader.read(symbol, modifiers)
            return
        try:
            {
                window.key.P: self._print_info,
                window.key.R: self._random,
                window.key.E: self._previous_random,
                window.key.D: self._set_duration,
                window.key.UP: self._increase_duration,
                window.key.DOWN: self._decrease_duration,
                window.key.LEFT: self._previous_slide,
                window.key.RIGHT: self._next_slide,
                window.key.SPACE: self._toggle_paused,
                window.key.I: self._jump,
                window.key.F: self.toggle_fullscreen,
                window.key.SLASH: self._search,
                window.key.ESCAPE: app.exit,
            }[symbol]()
        except KeyError:
            pass
<commit_msg>Some refactoring.<commit_after># Copyright (c) 2011 Bogdan Popa
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
from pyglet import app, gl, text, window
from string import ascii_letters, digits

from ymage.helpers import reschedule, reschedule_once
from ymage.slideshow import Slideshow

class Printer(text.Label):
    def __init__(self):
        super(Printer, self).__init__(
            "", "Arial", 12, False, False,
            (0, 255, 0, 255), 10, 10
        )

    def _print(self, message, duration=5):
        self.text = message
        reschedule_once(self.clear, duration)

    def clear(self, dt=None):
        self.text = ""

class Reader(object):
    def __init__(self, printer):
        self.is_reading = False
        self.callback = None
        self.printer = printer
        self.message = ""
        self._input = ""

    def read(self, symbol, modifiers):
        self.is_reading = True
        if symbol == window.key.RETURN:
            self.printer.clear()
            self.toggle_reading()
            self.callback(self._input)
            return
        elif symbol == window.key.ESCAPE:
            self.printer.clear()
            self.toggle_reading()
            return
        elif symbol == window.key.BACKSPACE:
            self._input = self._input[:-1]
        try:
            self._input += {
                window.key.PERIOD: ".",
                window.key.MINUS: "-",
                window.key.SPACE: " ",
            }[symbol]
        except KeyError:
            pass
        representation = window.key.symbol_string(symbol)
        representation = representation.replace("_", "")
        if representation in ascii_letters + digits:
            if modifiers & window.key.MOD_SHIFT:
                self._input += representation
            else:
                self._input += representation.lower()
        self.printer._print(
            "%s: %s_" % (self.message, self._input)
        )

    def start_reading(self, message, callback):
        self.callback = callback
        self.message = message
        self._input = ""
        self.printer._print(
            "%s: _" % self.message
        )
        self.toggle_reading()

    def toggle_reading(self):
        self.is_reading = not self.is_reading

class Window(window.Window):
    def __init__(self, options):
        super(Window, self).__init__(
            800, 600, "Loading...", True, None, False
        )
        self.slideshow = Slideshow(options)
        self.options = options
        self.printer = Printer()
        self.reader = Reader(self.printer)
        self.setup()

    def setup(self):
        if not self.options.windowed:
            self.toggle_fullscreen()

    def toggle_fullscreen(self):
        self.activate()
        self.set_mouse_visible(self.fullscreen)
        self.set_fullscreen(not self.fullscreen)

    def update_caption(self):
        self.set_caption("ymage [%d/%d]" % (
            self.slideshow.index + 1, len(self.slideshow.slides)
        ))

    # Actions
    def _set_duration(self):
        self.reader.start_reading(
            "Duration", self.slideshow.set_duration
        )

    def _increase_duration(self):
        self.printer._print("Duration: {}".format(
            self.slideshow.increase_duration()
        ))

    def _decrease_duration(self):
        self.printer._print("Duration: {}".format(
            self.slideshow.decrease_duration()
        ))

    def _previous_slide(self):
        self.slideshow.display(action="previous")

    def _next_slide(self):
        self.slideshow.display(action="next")

    def _toggle_paused(self):
        self.printer._print(
            self.slideshow.toggle_paused()
        )

    def _random(self):
        self.slideshow.display(action="random", previous=False)

    def _previous_random(self):
        self.slideshow.display(action="random", previous=True)

    def _print_info(self):
        self.printer._print("{} [{}/{}]".format(
            self.slideshow.get_current(),
            self.slideshow.index + 1,
            len(self.slideshow.slides)
        ))

    def _search(self, query=None):
        if query is None:
            self.reader.start_reading("Search for", self._search)
        else:
            self.slideshow.display(action="search", query=query)

    def _jump(self, index=None):
        if index is None:
            self.reader.start_reading("Slide", self._jump)
        else:
            self.slideshow.display(action="jump", index=index)

    def on_draw(self):
        self.clear()
        try:
            self.slideshow.draw_slide(self.width, self.height)
        except gl.lib.GLException:
            # In case one of the slides is corrupted
            # move on to the next
            self.slideshow.display()
        self.update_caption()
        self.printer.draw()

    def on_key_press(self, symbol, modifiers):
        if self.reader.is_reading:
            self.reader.read(symbol, modifiers)
            return
        try:
            {
                window.key.P: self._print_info,
                window.key.R: self._random,
                window.key.E: self._previous_random,
                window.key.D: self._set_duration,
                window.key.UP: self._increase_duration,
                window.key.DOWN: self._decrease_duration,
                window.key.LEFT: self._previous_slide,
                window.key.RIGHT: self._next_slide,
                window.key.SPACE: self._toggle_paused,
                window.key.I: self._jump,
                window.key.F: self.toggle_fullscreen,
                window.key.SLASH: self._search,
                window.key.ESCAPE: app.exit,
            }[symbol]()
        except KeyError:
            pass
<|endoftext|>


<commit_before>from collections import Mapping
from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, \
    create_string_buffer
from weakref import WeakValueDictionary

import numpy as np
from numpy.ctypeslib import as_array

from . import _dll
from .core import _FortranObjectWithID
from .error import _error_handler, AllocationError, InvalidIDError
from .material import Material

__all__ = ['Filter', 'AzimuthalFilter', 'CellFilter',
           'CellbornFilter', 'CellfromFilter', 'DistribcellFilter',
           'DelayedGroupFilter', 'EnergyFilter', 'EnergyoutFilter',
           'EnergyFunctionFilter', 'MaterialFilter', 'MeshFilter',
           'MuFilter', 'PolarFilter', 'SurfaceFilter',
           'UniverseFilter', 'filters']

# Tally functions
_dll.openmc_energy_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_double)), POINTER(c_int32)]
_dll.openmc_energy_filter_get_bins.restype = c_int
_dll.openmc_energy_filter_get_bins.errcheck = _error_handler
_dll.openmc_energy_filter_set_bins.argtypes = [c_int32, c_int32, POINTER(c_double)]
_dll.openmc_energy_filter_set_bins.restype = c_int
_dll.openmc_energy_filter_set_bins.errcheck = _error_handler
_dll.openmc_extend_filters.argtypes = [c_int32, POINTER(c_int32), POINTER(c_int32)]
_dll.openmc_extend_filters.restype = c_int
_dll.openmc_extend_filters.errcheck = _error_handler
_dll.openmc_filter_get_id.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_filter_get_id.restype = c_int
_dll.openmc_filter_get_id.errcheck = _error_handler
_dll.openmc_filter_get_type.argtypes = [c_int32, c_char_p]
_dll.openmc_filter_get_type.restype = c_int
_dll.openmc_filter_get_type.errcheck = _error_handler
_dll.openmc_filter_set_id.argtypes = [c_int32, c_int32]
_dll.openmc_filter_set_id.restype = c_int
_dll.openmc_filter_set_id.errcheck = _error_handler
_dll.openmc_filter_set_type.argtypes = [c_int32, c_char_p]
_dll.openmc_filter_set_type.restype = c_int
_dll.openmc_filter_set_type.errcheck = _error_handler
_dll.openmc_get_filter_index.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_get_filter_index.restype = c_int
_dll.openmc_get_filter_index.errcheck = _error_handler
_dll.openmc_material_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_int32)), POINTER(c_int32)]
_dll.openmc_material_filter_get_bins.restype = c_int
_dll.openmc_material_filter_get_bins.errcheck = _error_handler
_dll.openmc_material_filter_set_bins.argtypes = [c_int32, c_int32, POINTER(c_int32)]
_dll.openmc_material_filter_set_bins.restype = c_int
_dll.openmc_material_filter_set_bins.errcheck = _error_handler
_dll.openmc_mesh_filter_set_mesh.argtypes = [c_int32, c_int32]
_dll.openmc_mesh_filter_set_mesh.restype = c_int
_dll.openmc_mesh_filter_set_mesh.errcheck = _error_handler

class Filter(_FortranObjectWithID):
    __instances = WeakValueDictionary()

    def __new__(cls, filter_type, uid=None, new=True, index=None):
        mapping = filters
        if index is None:
            if new:
                # Determine ID to assign
                if uid is None:
                    try:
                        uid = max(mapping) + 1
                    except ValueError:
                        uid = 1
                else:
                    if uid in mapping:
                        raise AllocationError('A filter with ID={} has already '
                                              'been allocated.'.format(uid))

                index = c_int32()
                _dll.openmc_extend_filters(1, index, None)
                _dll.openmc_filter_set_type(index, filter_type)
                index = index.value
            else:
                index = mapping[uid]._index

        if index not in cls.__instances:
            instance = super(Filter, cls).__new__(cls)
            instance._index = index
            if uid is not None:
                instance.id = uid
            cls.__instances[index] = instance

        return cls.__instances[index]

    @property
    def id(self):
        filter_id = c_int32()
        _dll.openmc_filter_get_id(self._index, filter_id)
        return filter_id.value

    @id.setter
    def id(self, filter_id):
        _dll.openmc_filter_set_id(self._index, filter_id)

class EnergyFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(EnergyFilter, cls).__new__(cls, b'energy', uid, new, index)

    def __init__(self, bins=None, uid=None, new=True, index=None):
        super(EnergyFilter, self).__init__(uid, new, index)
        if bins is not None:
            self.bins = bins

    @property
    def bins(self):
        energies = POINTER(c_double)()
        n = c_int32()
        _dll.openmc_energy_filter_get_bins(self._index, energies, n)
        return as_array(energies, (n.value,))

    @bins.setter
    def bins(self, bins):
        # Get numpy array as a double*
        energies = np.asarray(bins)
        energies_p = energies.ctypes.data_as(POINTER(c_double))

        _dll.openmc_energy_filter_set_bins(
            self._index, len(energies), energies_p)

class EnergyoutFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(EnergyoutFilter, cls).__new__(cls, b'energyout',
                                                   uid, new, index)

class AzimuthalFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(AzimuthalFilter, cls).__new__(cls, b'azimuthal',
                                                   uid, new, index)

class CellFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(CellFilter, cls).__new__(cls, b'cell', uid, new, index)

class CellbornFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(CellbornFilter, cls).__new__(cls, b'cellborn', uid,
                                                  new, index)

class CellfromFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(CellfromFilter, cls).__new__(cls, b'cellfrom', uid,
                                                  new, index)

class DelayedGroupFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(DelayedGroupFilter, cls).__new__(cls, b'delayedgroup',
                                                      uid, new, index)

class DistribcellFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(DistribcellFilter, cls).__new__(cls, b'distribcell',
                                                     uid, new, index)

class EnergyFunctionFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(EnergyFunctionFilter, cls).__new__(
            cls, b'energyfunction', uid, new, index)

class MaterialFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(MaterialFilter, cls).__new__(cls, b'material',
                                                   uid, new, index)

    def __init__(self, bins=None, uid=None, new=True, index=None):
        super(MaterialFilter, self).__init__(uid, new, index)
        if bins is not None:
            self.bins = bins

    @property
    def bins(self):
        materials = POINTER(c_int32)()
        n = c_int32()
        _dll.openmc_material_filter_get_bins(self._index, materials, n)
        return [Material(index=materials[i]) for i in range(n.value)]

    @bins.setter
    def bins(self, materials):
        # Get material indices as int32_t[]
        n = len(materials)
        bins = (c_int32*n)(*(m._index for m in materials))

        _dll.openmc_material_filter_set_bins(self._index, n, bins)

class MeshFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(MeshFilter, cls).__new__(cls, b'mesh', uid, new, index)

class MuFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(MuFilter, cls).__new__(cls, b'mu', uid, new, index)

class PolarFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(PolarFilter, cls).__new__(cls, b'polar', uid, new, index)

class SurfaceFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(SurfaceFilter, cls).__new__(cls, b'surface', uid,
                                                 new, index)

class UniverseFilter(Filter):
    def __new__(cls, bins=None, uid=None, new=True, index=None):
        return super(UniverseFilter, cls).__new__(cls, b'universe', uid,
                                                  new, index)

_FILTER_TYPE_MAP = {
    'azimuthal': AzimuthalFilter,
    'cell': CellFilter,
    'cellborn': CellbornFilter,
    'cellfrom': CellfromFilter,
    'delayedgroup': DelayedGroupFilter,
    'distribcell': DistribcellFilter,
    'energy': EnergyFilter,
    'energyout': EnergyoutFilter,
    'energyfunction': EnergyFunctionFilter,
    'material': MaterialFilter,
    'mesh': MeshFilter,
    'mu': MuFilter,
    'polar': PolarFilter,
    'surface': SurfaceFilter,
    'universe': UniverseFilter,
}

def _get_filter(index):
    filter_type = create_string_buffer(20)
    _dll.openmc_filter_get_type(index, filter_type)
    filter_type = filter_type.value.decode()
    return _FILTER_TYPE_MAP[filter_type](index=index)

class _FilterMapping(Mapping):
    def __getitem__(self, key):
        index = c_int32()
        try:
            _dll.openmc_get_filter_index(key, index)
        except (AllocationError, InvalidIDError) as e:
            # __contains__ expects a KeyError to work correctly
            raise KeyError(str(e))
        return _get_filter(index.value)

    def __iter__(self):
        for i in range(len(self)):
            yield _get_filter(i + 1).id

    def __len__(self):
        return c_int32.in_dll(_dll, 'n_filters').value

    def __repr__(self):
        return repr(dict(self))

filters = _FilterMapping()
<commit_msg>Use class attribute in capi.Filter.__new__ to avoid the need to define __new__ in subclasses.<commit_after>from collections import Mapping
from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, \
    create_string_buffer
from weakref import WeakValueDictionary

import numpy as np
from numpy.ctypeslib import as_array

from . import _dll
from .core import _FortranObjectWithID
from .error import _error_handler, AllocationError, InvalidIDError
from .material import Material

__all__ = ['Filter', 'AzimuthalFilter', 'CellFilter',
           'CellbornFilter', 'CellfromFilter', 'DistribcellFilter',
           'DelayedGroupFilter', 'EnergyFilter', 'EnergyoutFilter',
           'EnergyFunctionFilter', 'MaterialFilter', 'MeshFilter',
           'MuFilter', 'PolarFilter', 'SurfaceFilter',
           'UniverseFilter', 'filters']

# Tally functions
_dll.openmc_energy_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_double)), POINTER(c_int32)]
_dll.openmc_energy_filter_get_bins.restype = c_int
_dll.openmc_energy_filter_get_bins.errcheck = _error_handler
_dll.openmc_energy_filter_set_bins.argtypes = [c_int32, c_int32, POINTER(c_double)]
_dll.openmc_energy_filter_set_bins.restype = c_int
_dll.openmc_energy_filter_set_bins.errcheck = _error_handler
_dll.openmc_extend_filters.argtypes = [c_int32, POINTER(c_int32), POINTER(c_int32)]
_dll.openmc_extend_filters.restype = c_int
_dll.openmc_extend_filters.errcheck = _error_handler
_dll.openmc_filter_get_id.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_filter_get_id.restype = c_int
_dll.openmc_filter_get_id.errcheck = _error_handler
_dll.openmc_filter_get_type.argtypes = [c_int32, c_char_p]
_dll.openmc_filter_get_type.restype = c_int
_dll.openmc_filter_get_type.errcheck = _error_handler
_dll.openmc_filter_set_id.argtypes = [c_int32, c_int32]
_dll.openmc_filter_set_id.restype = c_int
_dll.openmc_filter_set_id.errcheck = _error_handler
_dll.openmc_filter_set_type.argtypes = [c_int32, c_char_p]
_dll.openmc_filter_set_type.restype = c_int
_dll.openmc_filter_set_type.errcheck = _error_handler
_dll.openmc_get_filter_index.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_get_filter_index.restype = c_int
_dll.openmc_get_filter_index.errcheck = _error_handler
_dll.openmc_material_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_int32)), POINTER(c_int32)]
_dll.openmc_material_filter_get_bins.restype = c_int
_dll.openmc_material_filter_get_bins.errcheck = _error_handler
_dll.openmc_material_filter_set_bins.argtypes = [c_int32, c_int32, POINTER(c_int32)]
_dll.openmc_material_filter_set_bins.restype = c_int
_dll.openmc_material_filter_set_bins.errcheck = _error_handler
_dll.openmc_mesh_filter_set_mesh.argtypes = [c_int32, c_int32]
_dll.openmc_mesh_filter_set_mesh.restype = c_int
_dll.openmc_mesh_filter_set_mesh.errcheck = _error_handler

class Filter(_FortranObjectWithID):
    __instances = WeakValueDictionary()

    def __new__(cls, obj=None, uid=None, new=True, index=None):
        mapping = filters
        if index is None:
            if new:
                # Determine ID to assign
                if uid is None:
                    try:
                        uid = max(mapping) + 1
                    except ValueError:
                        uid = 1
                else:
                    if uid in mapping:
                        raise AllocationError('A filter with ID={} has already '
                                              'been allocated.'.format(uid))

                # Resize internal array
                index = c_int32()
                _dll.openmc_extend_filters(1, index, None)

                # Set the filter type -- note that the filter_type attribute
                # only exists on subclasses!
                _dll.openmc_filter_set_type(index, cls.filter_type.encode())
                index = index.value
            else:
                index = mapping[uid]._index

        if index not in cls.__instances:
            instance = super(Filter, cls).__new__(cls)
            instance._index = index
            if uid is not None:
                instance.id = uid
            cls.__instances[index] = instance

        return cls.__instances[index]

    @property
    def id(self):
        filter_id = c_int32()
        _dll.openmc_filter_get_id(self._index, filter_id)
        return filter_id.value

    @id.setter
    def id(self, filter_id):
        _dll.openmc_filter_set_id(self._index, filter_id)

class EnergyFilter(Filter):
    filter_type = 'energy'

    def __init__(self, bins=None, uid=None, new=True, index=None):
        super(EnergyFilter, self).__init__(uid, new, index)
        if bins is not None:
            self.bins = bins

    @property
    def bins(self):
        energies = POINTER(c_double)()
        n = c_int32()
        _dll.openmc_energy_filter_get_bins(self._index, energies, n)
        return as_array(energies, (n.value,))

    @bins.setter
    def bins(self, bins):
        # Get numpy array as a double*
        energies = np.asarray(bins)
        energies_p = energies.ctypes.data_as(POINTER(c_double))

        _dll.openmc_energy_filter_set_bins(
            self._index, len(energies), energies_p)

class EnergyoutFilter(Filter):
    filter_type = 'energyout'

class AzimuthalFilter(Filter):
    filter_type = 'azimuthal'

class CellFilter(Filter):
    filter_type = 'cell'

class CellbornFilter(Filter):
    filter_type = 'cellborn'

class CellfromFilter(Filter):
    filter_type = 'cellfrom'

class DelayedGroupFilter(Filter):
    filter_type = 'delayedgroup'

class DistribcellFilter(Filter):
    filter_type = 'distribcell'

class EnergyFunctionFilter(Filter):
    filter_type = 'energyfunction'

class MaterialFilter(Filter):
    filter_type = 'material'

    def __init__(self, bins=None, uid=None, new=True, index=None):
        super(MaterialFilter, self).__init__(uid, new, index)
        if bins is not None:
            self.bins = bins

    @property
    def bins(self):
        materials = POINTER(c_int32)()
        n = c_int32()
        _dll.openmc_material_filter_get_bins(self._index, materials, n)
        return [Material(index=materials[i]) for i in range(n.value)]

    @bins.setter
    def bins(self, materials):
        # Get material indices as int32_t[]
        n = len(materials)
        bins = (c_int32*n)(*(m._index for m in materials))

        _dll.openmc_material_filter_set_bins(self._index, n, bins)

class MeshFilter(Filter):
    filter_type = 'mesh'

class MuFilter(Filter):
    filter_type = 'mu'

class PolarFilter(Filter):
    filter_type = 'polar'

class SurfaceFilter(Filter):
    filter_type = 'surface'

class UniverseFilter(Filter):
    filter_type = 'universe'

_FILTER_TYPE_MAP = {
    'azimuthal': AzimuthalFilter,
    'cell': CellFilter,
    'cellborn': CellbornFilter,
    'cellfrom': CellfromFilter,
    'delayedgroup': DelayedGroupFilter,
    'distribcell': DistribcellFilter,
    'energy': EnergyFilter,
    'energyout': EnergyoutFilter,
    'energyfunction': EnergyFunctionFilter,
    'material': MaterialFilter,
    'mesh': MeshFilter,
    'mu': MuFilter,
    'polar': PolarFilter,
    'surface': SurfaceFilter,
    'universe': UniverseFilter,
}

def _get_filter(index):
    filter_type = create_string_buffer(20)
    _dll.openmc_filter_get_type(index, filter_type)
    filter_type = filter_type.value.decode()
    return _FILTER_TYPE_MAP[filter_type](index=index)

class _FilterMapping(Mapping):
    def __getitem__(self, key):
        index = c_int32()
        try:
            _dll.openmc_get_filter_index(key, index)
        except (AllocationError, InvalidIDError) as e:
            # __contains__ expects a KeyError to work correctly
            raise KeyError(str(e))
        return _get_filter(index.value)

    def __iter__(self):
        for i in range(len(self)):
            yield _get_filter(i + 1).id

    def __len__(self):
        return c_int32.in_dll(_dll, 'n_filters').value

    def __repr__(self):
        return repr(dict(self))

filters = _FilterMapping()
<|endoftext|>