<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
urlextract.py - file with definition of URLExtract class

.. Created on 2016-07-29
.. Licence MIT
.. codeauthor:: Jan Lipovský <janlipovsky@gmail.com>, janlipovsky.cz
"""

import os
import re
import string
import sys
import urllib.request
from datetime import datetime, timedelta
from urllib.error import URLError, HTTPError

import idna
import uritools

__VERSION__ = "0.3.1"  #: version of URLExtract class

class URLExtract:
    """
    Class for finding and extracting URLs from given string

    **Example**

    .. code-block:: python

        from urlextract import URLExtract

        extractor = URLExtract()
        urls = extractor.find_urls("Let's have URL janlipovsky.cz as an example.")
        print(urls) # prints: ['janlipovsky.cz']

    """
    # file name of cached list of TLDs downloaded from IANA
    cache_file_name = '.tlds'

    def __init__(self):
        """
        Initialize function for URLExtract class.
        Tries to get cached .tlds, if cached file does not exist it will try to download new list from IANNA
        and save it to users home directory.
        """
        # get directory for cached file
        dir_path = os.path.dirname(__file__)
        if not os.access(dir_path, os.W_OK):
            # get path to home dir
            dir_path = os.path.expanduser('~')

        # full path for cached file with list of TLDs
        self.tld_list_path = os.path.join(dir_path, self.cache_file_name)
        if not os.access(self.tld_list_path, os.F_OK):
            if not self._download_tlds_list():
                sys.exit(-1)

        # check if cached file is readable
        if not os.access(self.tld_list_path, os.R_OK):
            print("ERROR: Cached file is not readable for current user. ({})".format(self.tld_list_path))
            sys.exit(-2)

        # try to update cache file when cache is older than 7 days
        if not self.update_when_older(7):
            print("WARNING: Could not update file, using old version of TLDs list. ({})".format(self.tld_list_path))

        self.tlds = None
        self.tlds_re = None
        self._reload_tlds_from_file()

        self.hostname_re = re.compile("^([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])$")

        self.stop_chars = list(string.whitespace) + ['\"', '\'', '<', '>', ';', '@']
        # characters that are allowed to be right after TLD
        self.after_tld_chars = list(string.whitespace) + ['/', '\"', '\'', '<', '?', ':', '.', ',']

    def _reload_tlds_from_file(self):
        """
        Reloads TLDs from file and compile regexp.
        """
        # check if cached file is readable
        if not os.access(self.tld_list_path, os.R_OK):
            print("ERROR: Cached file is not readable for current user. ({})".format(self.tld_list_path))
        else:
            self.tlds = sorted(self._load_cached_tlds(), key=len, reverse=True)
            self.tlds_re = re.compile('|'.join([re.escape(tld) for tld in self.tlds]))

    def _download_tlds_list(self):
        """
        Function downloads list of TLDs from IANA 'https://data.iana.org/TLD/tlds-alpha-by-domain.txt'.

        :return: True if list was downloaded, False in case of an error
        :rtype: bool
        """
        url_list = 'https://data.iana.org/TLD/tlds-alpha-by-domain.txt'

        # check if we can write cache file
        if os.access(self.tld_list_path, os.F_OK) and not os.access(self.tld_list_path, os.W_OK):
            print("ERROR: Cache file is not writable for current user. ({})".format(self.tld_list_path))
            return False

        req = urllib.request.Request(url_list)
        req.add_header('User-Agent', 'Mozilla/5.0 (Windows NT 6.0; WOW64; rv:24.0) Gecko/20100101 Firefox/24.0')
        with open(self.tld_list_path, 'w') as ftld:
            try:
                with urllib.request.urlopen(req) as f:
                    page = f.read().decode('utf-8')
                    ftld.write(page)
            except HTTPError as e:
                print("ERROR: Can not download list ot TLDs. (HTTPError: {})".format(e.reason))
                return False
            except URLError as e:
                print("ERROR: Can not download list ot TLDs. (URLError: {})".format(e.reason))
                return False
        return True

    def _load_cached_tlds(self):
        """
        Loads TLDs from cached file to set.

        :return: Set of current TLDs
        :rtype: set
        """

        list_of_tlds = set()
        with open(self.tld_list_path, 'r') as f:
            for line in f:
                tld = line.strip().lower()
                # skip empty lines
                if len(tld) <= 0:
                    continue
                # skip comments
                if tld[0] == '#':
                    continue

                list_of_tlds.add("." + tld)
                list_of_tlds.add("." + idna.decode(tld))

        return list_of_tlds

    def _get_last_cachefile_modification(self):
        """
        Get last modification of cache file with tlds.

        :return: Date and time of last modification or None when file does not exist
        :rtype: datetime|None
        """

        try:
            mtime = os.path.getmtime(self.tld_list_path)
        except OSError:
            return None

        return datetime.fromtimestamp(mtime)

    def update(self):
        """
        Update tld list cache file.

        :return: True if update was successfull False otherwise
        :rtype: bool
        """

        if not self._download_tlds_list():
            return False

        self._reload_tlds_from_file()

        return True

    def update_when_older(self, days):
        """
        Update tld list cache file if the list is older than number of days given in parameter `days`.

        :param int days: number of days from last change
        :return: True if update was successfull, False otherwise
        :rtype: bool
        """

        last_cache = self._get_last_cachefile_modification()
        if last_cache is None:
            return False

        time_to_update = last_cache + timedelta(days=days)

        if datetime.now() >= time_to_update:
            return self.update()

        return True

    @staticmethod
    def get_version():
        """
        Returns version number.

        :return: version number
        :rtype: str
        """

        return __VERSION__

    def get_after_tld_chars(self):
        """
        Returns list of chars that are allowed after TLD

        :return: list of chars that are allowed after TLD
        :rtype: list
        """

        return self.after_tld_chars

    def set_after_tld_chars(self, after_tld_chars):
        """
        Set chars that are allowed after TLD.

        :param list after_tld_chars: list of characters
        """

        self.after_tld_chars = after_tld_chars

    def get_stop_chars(self):
        """
        Returns list of stop chars.

        :return: list of stop chars
        :rtype: list
        """

        return self.stop_chars

    def set_stop_chars(self, stop_chars):
        """
        Set stop characters used when determining end of URL.

        :param list stop_chars: list of characters
        """

        self.stop_chars = stop_chars

    def _complete_url(self, text, tld_pos):
        """
        Expand string in both sides to match whole URL.

        :param str text: text where we want to find URL
        :param int tld_pos: position of tld
        :return: returns URL
        :rtype: str
        """

        left_ok = True
        right_ok = True

        max_len = len(text) - 1
        end_pos = tld_pos
        start_pos = tld_pos
        while left_ok or right_ok:
            if left_ok:
                if start_pos <= 0:
                    left_ok = False
                else:
                    if text[start_pos - 1] not in self.stop_chars:
                        start_pos -= 1
                    else:
                        left_ok = False
            if right_ok:
                if end_pos >= max_len:
                    right_ok = False
                else:
                    if text[end_pos + 1] not in self.stop_chars:
                        end_pos += 1
                    else:
                        right_ok = False

        complete_url = text[start_pos:end_pos + 1].lstrip('/')
        if not self._is_domain_valid(complete_url):
            return ""

        return complete_url

    def _validate_tld_match(self, text, matched_tld, tld_pos):
        """
        Validate tld match - tells if at found position is really tld.

        :param str text: text where we want to find URLs
        :param str matched_tld: matched tld
        :param int tld_pos: position of matched tld
        :return: True if match is valid, False otherwise
        :rtype: bool
        """
        right_tld_pos = tld_pos + len(matched_tld)
        if len(text) > right_tld_pos:
            if text[right_tld_pos] in self.after_tld_chars:
                if tld_pos > 0 and text[tld_pos - 1] not in self.stop_chars:
                    return True
        else:
            if tld_pos > 0 and text[tld_pos - 1] not in self.stop_chars:
                return True

        return False

    def _is_domain_valid(self, url):
        """
        Checks if given URL has valid domain name (ignores subdomains)

        :param str url: text where we want to find URLs
        :return: True if URL is valid, False otherwise
        :rtype: bool

        >>> extractor = URLExtract()
        >>> extractor._is_domain_valid("janlipovsky.cz")
        True

        >>> extractor._is_domain_valid("https://janlipovsky.cz")
        True

        >>> extractor._is_domain_valid("invalid.cz.")
        False

        >>> extractor._is_domain_valid("in.v_alid.cz")
        False

        >>> extractor._is_domain_valid("-is.valid.cz")
        True

        >>> extractor._is_domain_valid("not.valid-.cz")
        False
        """

        if len(url) <= 0:
            return False

        scheme_pos = url.find('://')
        if scheme_pos != -1:
            url = url[scheme_pos+3:]

        url = 'http://'+url

        url_parts = uritools.urisplit(url)
        # <scheme>://<authority>/<path>?<query>#<fragment>
        host = url_parts.host
        if not host:
            return False

        host_parts = host.split('.')
        if len(host_parts) <= 1:
            return False

        tld = '.'+host_parts[-1]
        if tld not in self.tlds:
            return False

        top = host_parts[-2]

        if self.hostname_re.match(top) is None:
            return False

        return True

    def find_urls(self, text, only_unique=False):
        """
        Find all URLs in given text.

        >>> extractor = URLExtract()
        >>> extractor.find_urls("Let's have URL http://janlipovsky.cz as an example.")
        ['http://janlipovsky.cz']

        >>> extractor.find_urls("Let's have text without URLs.")
        []

        >>> extractor.find_urls("Get unique URL from: http://janlipovsky.cz http://janlipovsky.cz", True)
        ['http://janlipovsky.cz']

        :param str text: text where we want to find URLs
        :param bool only_unique: return only unique URLs
        :return: list of URLs found in text
        :rtype: list
        """
        urls = []
        tld_pos = 0
        matched_tlds = self.tlds_re.findall(text)

        for tld in matched_tlds:
            tmp_text = text[tld_pos:]
            offset = tld_pos
            tld_pos = tmp_text.find(tld)
            if tld_pos != -1 and self._validate_tld_match(text, tld, offset + tld_pos):
                urls.append(self._complete_url(text, offset + tld_pos))
            tld_pos += len(tld) + offset

        return urls if not only_unique else list(set(urls))
<commit_msg>Adding test case for bugfix<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
urlextract.py - file with definition of URLExtract class

.. Created on 2016-07-29
.. Licence MIT
.. codeauthor:: Jan Lipovský <janlipovsky@gmail.com>, janlipovsky.cz
"""

import os
import re
import string
import sys
import urllib.request
from datetime import datetime, timedelta
from urllib.error import URLError, HTTPError

import idna
import uritools

__VERSION__ = "0.3.2"  #: version of URLExtract class

class URLExtract:
    """
    Class for finding and extracting URLs from given string

    **Example**

    .. code-block:: python

        from urlextract import URLExtract

        extractor = URLExtract()
        urls = extractor.find_urls("Let's have URL janlipovsky.cz as an example.")
        print(urls) # prints: ['janlipovsky.cz']

    """
    # file name of cached list of TLDs downloaded from IANA
    cache_file_name = '.tlds'

    def __init__(self):
        """
        Initialize function for URLExtract class.
        Tries to get cached .tlds, if cached file does not exist it will try to download new list from IANNA
        and save it to users home directory.
        """
        # get directory for cached file
        dir_path = os.path.dirname(__file__)
        if not os.access(dir_path, os.W_OK):
            # get path to home dir
            dir_path = os.path.expanduser('~')

        # full path for cached file with list of TLDs
        self.tld_list_path = os.path.join(dir_path, self.cache_file_name)
        if not os.access(self.tld_list_path, os.F_OK):
            if not self._download_tlds_list():
                sys.exit(-1)

        # check if cached file is readable
        if not os.access(self.tld_list_path, os.R_OK):
            print("ERROR: Cached file is not readable for current user. ({})".format(self.tld_list_path))
            sys.exit(-2)

        # try to update cache file when cache is older than 7 days
        if not self.update_when_older(7):
            print("WARNING: Could not update file, using old version of TLDs list. ({})".format(self.tld_list_path))

        self.tlds = None
        self.tlds_re = None
        self._reload_tlds_from_file()

        self.hostname_re = re.compile("^([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])$")

        self.stop_chars = list(string.whitespace) + ['\"', '\'', '<', '>', ';', '@']
        # characters that are allowed to be right after TLD
        self.after_tld_chars = list(string.whitespace) + ['/', '\"', '\'', '<', '?', ':', '.', ',']

    def _reload_tlds_from_file(self):
        """
        Reloads TLDs from file and compile regexp.
        """
        # check if cached file is readable
        if not os.access(self.tld_list_path, os.R_OK):
            print("ERROR: Cached file is not readable for current user. ({})".format(self.tld_list_path))
        else:
            self.tlds = sorted(self._load_cached_tlds(), key=len, reverse=True)
            self.tlds_re = re.compile('|'.join([re.escape(tld) for tld in self.tlds]))

    def _download_tlds_list(self):
        """
        Function downloads list of TLDs from IANA 'https://data.iana.org/TLD/tlds-alpha-by-domain.txt'.

        :return: True if list was downloaded, False in case of an error
        :rtype: bool
        """
        url_list = 'https://data.iana.org/TLD/tlds-alpha-by-domain.txt'

        # check if we can write cache file
        if os.access(self.tld_list_path, os.F_OK) and not os.access(self.tld_list_path, os.W_OK):
            print("ERROR: Cache file is not writable for current user. ({})".format(self.tld_list_path))
            return False

        req = urllib.request.Request(url_list)
        req.add_header('User-Agent', 'Mozilla/5.0 (Windows NT 6.0; WOW64; rv:24.0) Gecko/20100101 Firefox/24.0')
        with open(self.tld_list_path, 'w') as ftld:
            try:
                with urllib.request.urlopen(req) as f:
                    page = f.read().decode('utf-8')
                    ftld.write(page)
            except HTTPError as e:
                print("ERROR: Can not download list ot TLDs. (HTTPError: {})".format(e.reason))
                return False
            except URLError as e:
                print("ERROR: Can not download list ot TLDs. (URLError: {})".format(e.reason))
                return False
        return True

    def _load_cached_tlds(self):
        """
        Loads TLDs from cached file to set.

        :return: Set of current TLDs
        :rtype: set
        """

        list_of_tlds = set()
        with open(self.tld_list_path, 'r') as f:
            for line in f:
                tld = line.strip().lower()
                # skip empty lines
                if len(tld) <= 0:
                    continue
                # skip comments
                if tld[0] == '#':
                    continue

                list_of_tlds.add("." + tld)
                list_of_tlds.add("." + idna.decode(tld))

        return list_of_tlds

    def _get_last_cachefile_modification(self):
        """
        Get last modification of cache file with tlds.

        :return: Date and time of last modification or None when file does not exist
        :rtype: datetime|None
        """

        try:
            mtime = os.path.getmtime(self.tld_list_path)
        except OSError:
            return None

        return datetime.fromtimestamp(mtime)

    def update(self):
        """
        Update tld list cache file.

        :return: True if update was successfull False otherwise
        :rtype: bool
        """

        if not self._download_tlds_list():
            return False

        self._reload_tlds_from_file()

        return True

    def update_when_older(self, days):
        """
        Update tld list cache file if the list is older than number of days given in parameter `days`.

        :param int days: number of days from last change
        :return: True if update was successfull, False otherwise
        :rtype: bool
        """

        last_cache = self._get_last_cachefile_modification()
        if last_cache is None:
            return False

        time_to_update = last_cache + timedelta(days=days)

        if datetime.now() >= time_to_update:
            return self.update()

        return True

    @staticmethod
    def get_version():
        """
        Returns version number.

        :return: version number
        :rtype: str
        """

        return __VERSION__

    def get_after_tld_chars(self):
        """
        Returns list of chars that are allowed after TLD

        :return: list of chars that are allowed after TLD
        :rtype: list
        """

        return self.after_tld_chars

    def set_after_tld_chars(self, after_tld_chars):
        """
        Set chars that are allowed after TLD.

        :param list after_tld_chars: list of characters
        """

        self.after_tld_chars = after_tld_chars

    def get_stop_chars(self):
        """
        Returns list of stop chars.

        :return: list of stop chars
        :rtype: list
        """

        return self.stop_chars

    def set_stop_chars(self, stop_chars):
        """
        Set stop characters used when determining end of URL.

        :param list stop_chars: list of characters
        """

        self.stop_chars = stop_chars

    def _complete_url(self, text, tld_pos):
        """
        Expand string in both sides to match whole URL.

        :param str text: text where we want to find URL
        :param int tld_pos: position of tld
        :return: returns URL
        :rtype: str
        """

        left_ok = True
        right_ok = True

        max_len = len(text) - 1
        end_pos = tld_pos
        start_pos = tld_pos
        while left_ok or right_ok:
            if left_ok:
                if start_pos <= 0:
                    left_ok = False
                else:
                    if text[start_pos - 1] not in self.stop_chars:
                        start_pos -= 1
                    else:
                        left_ok = False
            if right_ok:
                if end_pos >= max_len:
                    right_ok = False
                else:
                    if text[end_pos + 1] not in self.stop_chars:
                        end_pos += 1
                    else:
                        right_ok = False

        complete_url = text[start_pos:end_pos + 1].lstrip('/')
        if not self._is_domain_valid(complete_url):
            return ""

        return complete_url

    def _validate_tld_match(self, text, matched_tld, tld_pos):
        """
        Validate tld match - tells if at found position is really tld.

        :param str text: text where we want to find URLs
        :param str matched_tld: matched tld
        :param int tld_pos: position of matched tld
        :return: True if match is valid, False otherwise
        :rtype: bool
        """
        right_tld_pos = tld_pos + len(matched_tld)
        if len(text) > right_tld_pos:
            if text[right_tld_pos] in self.after_tld_chars:
                if tld_pos > 0 and text[tld_pos - 1] not in self.stop_chars:
                    return True
        else:
            if tld_pos > 0 and text[tld_pos - 1] not in self.stop_chars:
                return True

        return False

    def _is_domain_valid(self, url):
        """
        Checks if given URL has valid domain name (ignores subdomains)

        :param str url: text where we want to find URLs
        :return: True if URL is valid, False otherwise
        :rtype: bool

        >>> extractor = URLExtract()
        >>> extractor._is_domain_valid("janlipovsky.cz")
        True

        >>> extractor._is_domain_valid("https://janlipovsky.cz")
        True

        >>> extractor._is_domain_valid("invalid.cz.")
        False

        >>> extractor._is_domain_valid("in.v_alid.cz")
        False

        >>> extractor._is_domain_valid("-is.valid.cz")
        True

        >>> extractor._is_domain_valid("not.valid-.cz")
        False

        >>> extractor._is_domain_valid("http://blog/media/reflect.io.jpg")
        False
        """

        if len(url) <= 0:
            return False

        scheme_pos = url.find('://')
        if scheme_pos != -1:
            url = url[scheme_pos+3:]

        url = 'http://'+url

        url_parts = uritools.urisplit(url)
        # <scheme>://<authority>/<path>?<query>#<fragment>
        host = url_parts.host
        if not host:
            return False

        host_parts = host.split('.')
        if len(host_parts) <= 1:
            return False

        tld = '.'+host_parts[-1]
        if tld not in self.tlds:
            return False

        top = host_parts[-2]

        if self.hostname_re.match(top) is None:
            return False

        return True

    def find_urls(self, text, only_unique=False):
        """
        Find all URLs in given text.

        >>> extractor = URLExtract()
        >>> extractor.find_urls("Let's have URL http://janlipovsky.cz as an example.")
        ['http://janlipovsky.cz']

        >>> extractor.find_urls("Let's have text without URLs.")
        []

        >>> extractor.find_urls("Get unique URL from: http://janlipovsky.cz http://janlipovsky.cz", True)
        ['http://janlipovsky.cz']

        :param str text: text where we want to find URLs
        :param bool only_unique: return only unique URLs
        :return: list of URLs found in text
        :rtype: list
        """
        urls = []
        tld_pos = 0
        matched_tlds = self.tlds_re.findall(text)

        for tld in matched_tlds:
            tmp_text = text[tld_pos:]
            offset = tld_pos
            tld_pos = tmp_text.find(tld)
            if tld_pos != -1 and self._validate_tld_match(text, tld, offset + tld_pos):
                urls.append(self._complete_url(text, offset + tld_pos))
            tld_pos += len(tld) + offset

        return urls if not only_unique else list(set(urls))
<|endoftext|>


<commit_before>#  Copyright 2014 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
from Rammbock import logger

from Rammbock.binary_tools import to_bin

class MessageStream(object):

    def __init__(self, stream, protocol):
        self._cache = []
        self._stream = stream
        self._protocol = protocol

    def get(self, message_template, timeout=None, header_filter=None):
        header_fields = message_template.header_parameters
        logger.trace("Get message with params %s" % header_fields)
        msg = self._get_from_cache(message_template, header_fields, header_filter)
        if msg:
            logger.trace("Cache hit. Cache currently has %s messages" % len(self._cache))
            return msg
        while True:
            header, pdu_bytes = self._protocol.read(self._stream, timeout=timeout)
            if self._matches(header, header_fields, header_filter):
                return self._to_msg(message_template, header, pdu_bytes)
            self._cache.append((header, pdu_bytes))

    def _get_from_cache(self, template, fields, header_filter):
        for index in range(len(self._cache)):
            header, pdu = self._cache[index]
            if self._matches(header, fields, header_filter):
                self._cache.pop(index)
                return self._to_msg(template, header, pdu)
        return None

    def _to_msg(self, template, header, pdu_bytes):
        if template.only_header:
            return header
        msg = template.decode(pdu_bytes, parent=header)
        msg._add_header(header)
        return msg

    def _matches(self, header, fields, header_filter):
        if header_filter:
            if header_filter not in fields:
                raise AssertionError('Trying to filter messages by header field %s, but no value has been set for %s' %
                                     (header_filter, header_filter))
            if header[header_filter].bytes != to_bin(fields[header_filter]):
                return False
        return True

    def empty(self):
        self._cache = []
        self._stream.empty()

    def get_messages_count_in_cache(self):
        self._fill_cache()
        for msg in self._cache:
            logger.info(msg)
        return len(self._cache)

    def _fill_cache(self):
        try:
            while True:
                header, pdu_bytes = self._protocol.read(self._stream, timeout=0.2)
                self._cache.append((header, pdu_bytes))
        except:
            pass
<commit_msg>Added support for asynchronous messages<commit_after>#  Copyright 2014 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
from Rammbock import logger

from Rammbock.binary_tools import to_bin

class MessageStream(object):

    def __init__(self, stream, protocol):
        self._cache = []
        self._stream = stream
        self._protocol = protocol
        self._handlers = []

    def set_handler(self, msg_template, handler_func):
        self._handlers.append((msg_template, handler_func))

    def get(self, message_template, timeout=None, header_filter=None):
        header_fields = message_template.header_parameters
        logger.trace("Get message with params %s" % header_fields)
        msg = self._get_from_cache(message_template, header_fields, header_filter)
        if msg:
            logger.trace("Cache hit. Cache currently has %s messages" % len(self._cache))
            return msg
        while True:
            header, pdu_bytes = self._protocol.read(self._stream, timeout=timeout)
            if self._matches(header, header_fields, header_filter):
                return self._to_msg(message_template, header, pdu_bytes)
            else:
                self._match_or_cache(header, pdu_bytes, header_filter)

    def _match_or_cache(self, header, pdu_bytes, header_filter):
        for template, func in self._handlers:
            if self._matches(header, template.header_parameters, header_filter):
                msg_to_be_sent = self._to_msg(template, header, pdu_bytes)
                self._get_call_handler(func)(self, msg_to_be_sent)
                return  #FIXME Make test that handled messages dont go to cache
        self._cache.append((header, pdu_bytes))

    def _get_call_handler(self,handler_name):
        module, function = handler_name.split('.')
        mod = __import__(module)
        return getattr(mod, function)

    def _get_from_cache(self, template, fields, header_filter):
        for index in range(len(self._cache)):
            header, pdu = self._cache[index]
            if self._matches(header, fields, header_filter):
                self._cache.pop(index)
                return self._to_msg(template, header, pdu)
        return None

    def _to_msg(self, template, header, pdu_bytes):
        if template.only_header:
            return header
        msg = template.decode(pdu_bytes, parent=header)
        msg._add_header(header)
        return msg

    def _matches(self, header, fields, header_filter):
        # FIXME: Matching should not be assuming matching string presentation
        if header_filter:
            if header_filter not in fields:
                raise AssertionError('Trying to filter messages by header field %s, but no value has been set for %s' %
                                     (header_filter, header_filter))
            if header[header_filter].bytes != to_bin(fields[header_filter]):
                return False
        return True

    def empty(self):
        self._cache = []
        self._stream.empty()

    def get_messages_count_in_cache(self):
        self._fill_cache()
        for msg in self._cache:
            logger.info(msg)
        return len(self._cache)

    def _fill_cache(self):
        try:
            while True:
                header, pdu_bytes = self._protocol.read(self._stream, timeout=0.2)
                self._cache.append((header, pdu_bytes))
        except:
            pass
<|endoftext|>


<commit_before>#coding:utf-8
import os
import sys
import fcntl
import time
import re
from subprocess import Popen, PIPE

openpgp_trust = {"-": "Trust not calculated",
                 "o": "Unknown trust",
                 "q": "Undefined trust",
                 "n": "Never trust",
                 "m": "Marginally trust",
                 "f": "Full trust",
                 "u": "Ultimate trust",
                 "e": "Expired key, not trusted",
                 "r": "Revoked key, not trusted",
                 "d": "Disabled key, not trusted",  # Deprecated flag.
                }

openpgp_algorithms = {1: "RSA",
                      2: "RSA (encrypt only)",
                      3: "RSA (sign only)",
                      16: "Elgamal (encrypt only)",
                      17: "DSA",
                      20: "Elgamal (encrypt/sign) [COMPROMISED]",
                     }
# For details on type 20 compromisation, see http://lists.gnupg.org/pipermail/gnupg-announce/2003q4/000160.html

class GnuPG:
    """
    Wrap GnuPG and make all functionality feel Pythonic.
    """

    def __init__(self):
        self.gpgbinary = 'gpg'
        self.passphrase = None
        self.rw = ["r", "w"]
        self.fds = {"passphrase": True, "command": True, "logger": False, "status": False}
        self.handles = {}
        self.pipes = {}
        self.needed_fds = ["stdin", "stdout", "stderr"]
        self.errors = []

    def default_errorhandler(self, *error):
        if error != "":
            self.errors.append(error)
        return True

    def default_output(self, output):
        return output

    def parse_verify(self, output, *args):
        lines = output.split("\n")
        sig = {"datetime": "",
               "status": "",
               "keyid": "",
               "signer": "",
               "ok": None,
               "version": "",
               "hash": ""
              }

        if "no valid OpenPGP data found" in lines[0]:
            sig["ok"] = False
            sig["status"] = lines[1][4:]
        elif False:
            pass

        return sig

    def parse_keylist(self, keylist, *args):
        """
        >>> g = GnuPG()
        >>> v = g.parse_keylist("pub:u:4096:1:D5DC2A79C2E4AE92:2010-12-30:::u:Smari McCarthy <smari@immi.is>::scESC:\\nsub:u:4096:1:13E0BB42176BA0AC:2010-12-30::::::e:")
        >>> v.has_key("D5DC2A79C2E4AE92")
        True
        >>> v["D5DC2A79C2E4AE92"]["size"]
        4096
        >>> v["D5DC2A79C2E4AE92"]["creation-date"]
        '2010-12-30'
        >>> v["D5DC2A79C2E4AE92"]["algorithm"]
        1
        >>> v["D5DC2A79C2E4AE92"]["subkeys"][0]["algorithm"]
        1
        """
        keys = {}
        curkey = None

        def parse_pubkey(line, curkey, keys):
            keys[line[4]] = {
                "size": int(line[2]),
                "creation-date": line[5],
                "uids": [],
                "subkeys": [],
                "signatures": [],
                "trust": line[1],
                "algorithm": int(line[3])
            }
            if line[6] != "":
                keys[line[4]]["revocation-date"] = line[5]
            curkey = line[4]
            curkey, keys = parse_uid(line, curkey, keys)
            return (curkey, keys)

        def parse_subkey(line, curkey, keys):
            subkey = {"id": line[4], "size": int(line[2]), "creation-date": line[5], "algorithm": int(line[3])}
            if line[0] == "ssb":
                subkey["secret"] = True
            keys[curkey]["subkeys"].append(subkey)
            return (curkey, keys)

        def parse_fingerprint(line, curkey, keys):
            keys[curkey]["fingerprint"] = line[9]
            return (curkey, keys)

        def parse_userattribute(line, curkey, keys):
            # TODO: We are currently ignoring user attributes as not useful.
            #       We may at some point want to use --attribute-fd and read
            #       in user photos and such?
            return (curkey, keys)

        def parse_privkey(line, curkey, keys):
            curkey, keys = parse_pubkey(line, curkey, keys)
            return (curkey, keys)

        def parse_uid(line, curkey, keys):
            matches = re.match("([^\(\)]*)( \((.*)\)){0,1} \<(.*)\>", line[9])
            if matches:
                email = matches.groups(0)[3]
                comment = matches.groups(0)[2] or ""
                name = matches.groups(0)[0]
            else:
                email = line[9]
                name = ""
                comment = ""
            keys[curkey]["uids"].append({"email": email, "name": name, "comment": comment, "creation-date": line[5] })
            return (curkey, keys)

        def parse_trust(line, curkey, keys):
            # TODO: We are currently ignoring commentary from the Trust DB.
            return (curkey, keys)

        def parse_signature(line, curkey, keys):
            sig = {"signer": line[9], "signature-date": line[5], "keyid": line[4], "trust": line[10], "algorithm": line[4]}

            keys[curkey]["signatures"].append(sig)
            return (curkey, keys)

        def parse_unknown(line, curkey, keys):
            print "Unknown line with code '%s'" % line[0]
            return (curkey, keys)

        def parse_none(line, curkey, keys):
            return (curkey, keys)

        disp = {"pub": parse_pubkey,
                "sub": parse_subkey,
                "ssb": parse_subkey,
                "fpr": parse_fingerprint,
                "uat": parse_userattribute,
                "sec": parse_privkey,
                "tru": parse_trust,
                "sig": parse_signature,
                "uid": parse_uid,
                "gpg": parse_none,
               }

        lines = keylist.split("\n")
        for line in lines:
            if line == "":
                continue
            parms = line.split(":")
            r = disp.get(parms[0], parse_unknown)
            curkey, keys = r(parms, curkey, keys)

        return keys

    def emptycallbackmap():
        """Utility function for people who are confused about what callbacks exist."""
        return dict([[x, []] for x in self.needed_fds])

    def run(self, args=[], callbacks={}, output=None, debug=False):
        """
        >>> g = GnuPG()
        >>> g.run(["--list-keys"])[0]
        0
        """
        self.pipes = {}
        args.insert(0, self.gpgbinary)
        args.append("--with-colons")
        args.append("--verbose")
        args.append("--enable-progress-filter")
        if debug: print "Needed FDs: ", self.needed_fds
        for fd in self.fds.keys():
            if fd not in self.needed_fds:
                if debug: print "Don't need %s" % fd
                continue
            if debug: print "Opening fd %s" % fd
            self.pipes[fd] = os.pipe()
            args.append("--%s-fd=%d" % (fd, self.pipes[fd][not self.fds[fd]]))
            self.handles[fd] = os.fdopen(self.pipes[fd][self.fds[fd]], self.rw[self.fds[fd]])
            fcntl.fcntl(self.handles[fd], fcntl.F_SETFD, 0)   # Stay open after execing
            # Make the descriptor non-blocking
            flags = fcntl.fcntl(self.handles[fd], fcntl.F_GETFL)
            fcntl.fcntl(self.handles[fd], fcntl.F_SETFL, flags | os.O_NONBLOCK)

        if debug: print "Running gpg with %s" % args

        proc = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE)
        self.handles["stdout"] = proc.stdout
        self.handles["stderr"] = proc.stderr
        self.handles["stdin"] = proc.stdin
        if self.passphrase:
            self.handles["passphrase"].write(self.passphrase)
            self.handles["passphrase"].close()

        if output:
            self.handles["stdin"].write(output)
            self.handles["stdin"].close()

        retvals = {}
        while True:
            r = proc.poll()
            for fd in self.needed_fds:
                if fd in ("stdin", "passphrase"):
                    continue
                if debug: print "Reading %s" % fd
                buf = self.handles[fd].read()
                if not callbacks.has_key(fd):
                    continue
                if not retvals.has_key(fd):
                    retvals[fd] = []
                if buf != "":
                    if type(callbacks[fd]) == list:
                        for cb in callbacks[fd]:
                            retvals[fd].append(cb(buf))
                    else:
                        retvals[fd].append(callbacks[fd](buf))

            if proc.returncode is not None:
                break

        return proc.returncode, retvals

    def list_keys(self):
        """
        >>> g = GnuPG()
        >>> g.list_keys()[0]
        0
        """
        retvals = self.run(["--list-keys", "--fingerprint"], callbacks={"stdout": self.parse_keylist})
        return retvals

    def list_sigs(self):
        self.run(["--list-sigs", "--fingerprint"], callbacks={"stdout": self.parse_keylist})
        return retvals

    def list_secret_keys(self):
        """
        >>> g = GnuPG()
        >>> g.list_secret_keys()[0]
        0
        """
        retvals = self.run(["--list-secret-keys", "--fingerprint"], callbacks={"stdout": self.parse_keylist})
        return retvals

    def encrypt(self, data, to=[], armor=True):
        """
        >>> g = GnuPG()
        >>> g.encrypt("Hello, World", to=["smari@mailpile.is"])[0]
        0
        """
        action = ["--encrypt"]
        if armor:
            action.append("--armor")
        for r in to:
            action.append("--recipient")
            action.append(r)
        retvals = self.run(action, callbacks={"stdout": self.default_output}, output=data)
        return retvals[0], retvals[1]["stdout"][0]

    def decrypt(self, data, passphrase=None):
        """
        Note that this test will fail if you don't replace the recipient with one whose key you control.
        >>> g = GnuPG()
        >>> ct = g.encrypt("Hello, World", to=["smari@mailpile.is"])[1]
        >>> g.decrypt(ct)[1]
        'Hello, World'
        """
        if passphrase:
            self.passphrase = passphrase
        action = ["--decrypt"]
        retvals = self.run(action, callbacks={"stdout": self.default_output}, output=data)
        return retvals[0], retvals[1]["stdout"][0]

    def sign(self, data, _from=None, armor=True, detatch=True, clearsign=False, passphrase=None):
        """
        >>> g = GnuPG()
        >>> g.sign("Hello, World", _from="smari@mailpile.is")[0]
        0
        """
        if passphrase:
            self.passphrase = passphrase
        if detatch and not clearsign:
            action = ["--detach-sign"]
        elif clearsign:
            action = ["--clearsign"]
        else:
            action = ["--sign"]
        if armor:
            action.append("--armor")
        if _from:
            action.append("--local-user")
            action.append(_from)

        retvals = self.run(action, callbacks={"stdout": self.default_output}, output=data)
        return retvals[0], retvals[1]["stdout"][0]

    def verify(self, data, signature=None):
        """
        >>> g = GnuPG()
        >>> s = g.sign("Hello, World", _from="smari@mailpile.is", clearsign=True)[1]
        >>> g.verify(s)
        """
        print "BEFORE: %s" % self.needed_fds
        self.needed_fds.append("status")
        print "AFTER: %s" % self.needed_fds
        retvals = self.run(["--verify"], callbacks={"stderr": self.parse_verify, "status": self.default_output}, output=data, debug=True)
        self.needed_fds.remove("status")

        return retvals[0], retvals[1]["stderr"]

    def sign_encrypt(self, data, _from=None, to=[], armor=True, detatch=True, clearsign=False):
        retval, signblock = self.sign(data, _from=_from, armor=armor, detatch=detatch, clearsign=clearsign)
        if detatch:
            # TODO: Deal with detached signature.
            retval, cryptblock = self.encrypt(data, to=to, armor=armor)
        else:
            retval, cryptblock = self.encrypt(signblock, to=to, armor=armor)

        return cryptblock

if __name__ == "__main__":
    g = GnuPG()
    g.verify("Foo")
    import doctest
    t = doctest.testmod()
    if t.failed == 0:
        print "GPG Interface: All %d tests successful" % (t.attempted)
    else:
        print "GPG Interface: %d out of %d tests failed" % (t.failed, t.attempted)
<commit_msg>Adding status-fd status code management.<commit_after>#coding:utf-8
import os
import sys
import fcntl
import time
import re
from subprocess import Popen, PIPE

openpgp_trust = {"-": "Trust not calculated",
                 "o": "Unknown trust",
                 "q": "Undefined trust",
                 "n": "Never trust",
                 "m": "Marginally trust",
                 "f": "Full trust",
                 "u": "Ultimate trust",
                 "e": "Expired key, not trusted",
                 "r": "Revoked key, not trusted",
                 "d": "Disabled key, not trusted",  # Deprecated flag.
                }

openpgp_algorithms = {1: "RSA",
                      2: "RSA (encrypt only)",
                      3: "RSA (sign only)",
                      16: "Elgamal (encrypt only)",
                      17: "DSA",
                      20: "Elgamal (encrypt/sign) [COMPROMISED]",
                     }
# For details on type 20 compromisation, see
# http://lists.gnupg.org/pipermail/gnupg-announce/2003q4/000160.html

# These are detailed in the GnuPG source under doc/DETAILS.
status_messages = {
    "ENTER": [],
    "LEAVE": [],
    "ABORT": [],
    "NEWSIG": [],
    "GOODSIG": ["long_keyid_or_fpr", "username"],
    "KEYEXPIRED": ["expire_timestamp"],
    "KEYREVOKED": [],
    "BADSIG": ["long_keyid_or_fpr", "username"],
    "ERRSIG": ["long_keyid_or_fpr", "pubkey_algo", "hash_algo", "sig_class",
               "timestamp", "rc"],
    "BADARMOR": [],
    "TRUST_UNDEFINED": ["error_token"],
    "TRUST_NEVER": ["error_token"],
    "TRUST_MARGINAL": ["zero", "validation_model"],
    "TRUST_FULLY": ["zero", "validation_model"],
    "TRUST_ULTIMATE": ["zero", "validation_model"],
    "GET_BOOL": [],
    "GET_LINE": [],
    "GET_HIDDEN": [],
    "GOT_IT": [],
    "SHM_INFO": [],
    "SHM_GET": [],
    "SHM_GET_BOOL": [],
    "SHM_GET_HIDDEN": [],
    "NEED_PASSPHRASE": ["long_main_keyid", "long_keyid",
                        "keytype", "keylength"],
    "VALIDSIG": ["fingerprint", "sig_creation_date", "sig_timestamp",
                 "expire_timestamp","sig_version", "reserved", "pubkey_algo",
                 "hash_algo", "sig_class", "primary_key_fpr"],
    "SIG_ID": ["radix64_string", "sig_creation_date", "sig_timestamp"],
    "ENC_TO": ["long_keyid", "keytype", "keylength"],
    "NODATA": ["what"],
    "BAD_PASSPHRASE": ["long_keyid"],
    "NO_PUBKEY": ["long_keyid"],
    "NO_SECKEY": ["long_keyid"],
    "NEED_PASSPHRASE_SYM": ["cipher_algo", "s2k_mode", "s2k_hash"],
    "NEED_PASSPHRASE_PIN": ["card_type", "chvno", "serialno"],
    "DECRYPTION_FAILED": [],
    "DECRYPTION_OKAY": [],
    "MISSING_PASSPHRASE": [],
    "GOOD_PASSPHRASE": [],
    "GOODMDC": [],
    "BADMDC": [],
    "ERRMDC": [],
    "IMPORTED": ["long keyid", "username"],
    "IMPORT_OK": ["reason", "fingerprint"],
    "IMPORT_PROBLEM": ["reason", "fingerprint"],
    "IMPORT_CHECK": [],
    "IMPORT_RES": ["count", "no_user_id", "imported", "imported_rsa",
                   "unchanged", "n_uids", "n_subk", "n_sigs", "n_revoc",
                   "sec_read", "sec_imported", "sec_dups", "skipped_new_keys",
                   "not_imported"],
    "FILE_START": ["what", "filename"],
    "FILE_DONE": [],
    "FILE_ERROR": [],
    "BEGIN_DECRYPTION": ["mdc_method", "sym_algo"],
    "END_DECRYPTION": [],
    "BEGIN_ENCRYPTION": [],
    "END_ENCRYPTION": [],
    "DELETE_PROBLEM": ["reason_code"],
    "PROGRESS": ["what", "char", "cur", "total"],
    "SIG_CREATED": ["type" "pubkey algo", "hash algo", "class",
                    "timestamp", "key fpr"],
    "SESSION_KEY": ["algo:hexdigits"],
    "NOTATION_NAME" : ["name"],
    "NOTATION_DATA" : ["string"],
    "POLICY_URL" : ["string"],
    "BEGIN_STREAM": [],
    "END_STREAM": [],
    "KEY_CREATED": ["type", "fingerprint", "handle"],
    "KEY_NOT_CREATED": ["handle"],
    "USERID_HINT": ["long main keyid", "string"],
    "UNEXPECTED": ["what"],
    "INV_RECP": ["reason", "requested_recipient"],
    "INV_SGNR": ["reason", "requested_sender"],
    "NO_RECP": ["reserved"],
    "NO_SGNR": ["reserved"],
    "ALREADY_SIGNED": ["long-keyid"], # Experimental, may disappear
    "SIGEXPIRED": [], # Deprecated but may crop up; keyexpired overrides
    "TRUNCATED": ["maxno"],
    "EXPSIG": ["long_keyid_or_fpr", "username"],
    "EXPKEYSIG": ["long_keyid_or_fpr", "username"],
    "REVKEYSIG": ["long_keyid_or_fpr", "username"],
    "ATTRIBUTE": ["fpr", "octets", "type", "index",
                  "count", "timestamp", "expiredate", "flags"],
    "CARDCTRL": ["what", "serialno"],
    "PLAINTEXT": ["format", "timestamp", "filename"],
    "PLAINTEXT_LENGTH": ["length"],
    "SIG_SUBPACKET": ["type", "flags", "len", "data"],
    "SC_OP_SUCCESS": ["code"],
    "SC_OP_FAILURE": ["code"],
    "BACKUP_KEY_CREATED": ["fingerprint", "fname"],
    "PKA_TRUST_BAD": ["unknown"],
    "PKA_TRUST_GOOD": ["unknown"],
    "BEGIN_SIGNING": [],
    "ERROR": ["error location", "error code", "more"],
    "MOUNTPOINT": ["mdc_method", "sym_algo"],
    "SUCCESS": ["location"],
    "DECRYPTION_INFO": [],
}

class GnuPG:
    """
    Wrap GnuPG and make all functionality feel Pythonic.
    """

    def __init__(self):
        self.gpgbinary = 'gpg'
        self.passphrase = None
        self.rw = ["r", "w"]
        self.fds = {"passphrase": True, "command": True, "logger": False, "status": False}
        self.handles = {}
        self.pipes = {}
        self.needed_fds = ["stdin", "stdout", "stderr"]
        self.errors = []

    def default_errorhandler(self, *error):
        if error != "":
            self.errors.append(error)
        return True

    def default_output(self, output):
        return output

    def parse_status(self, output, *args):
        status = []
        lines = output.split("\n")
        for line in lines:
            line = line.replace("[GNUPG:] ", "")
            if line == "":
                continue
            elems = line.split(" ")
            callback_kwargs = dict(zip(status_messages, elems[1:]))
            if self.statuscallbacks.has_key(elems[0]):
                for callback in self.statuscallbacks[elems[0]]:
                    callback(*kwargs)
            status.append(elems)

        return status

    def parse_verify(self, output, *args):
        lines = output.split("\n")
        sig = {"datetime": "",
               "status": "",
               "keyid": "",
               "signer": "",
               "ok": None,
               "version": "",
               "hash": ""
              }

        if "no valid OpenPGP data found" in lines[0]:
            sig["ok"] = False
            sig["status"] = lines[1][4:]
        elif False:
            pass

        return sig

    def parse_keylist(self, keylist, *args):
        """
        >>> g = GnuPG()
        >>> v = g.parse_keylist("pub:u:4096:1:D5DC2A79C2E4AE92:2010-12-30:::u:Smari McCarthy <smari@immi.is>::scESC:\\nsub:u:4096:1:13E0BB42176BA0AC:2010-12-30::::::e:")
        >>> v.has_key("D5DC2A79C2E4AE92")
        True
        >>> v["D5DC2A79C2E4AE92"]["size"]
        4096
        >>> v["D5DC2A79C2E4AE92"]["creation-date"]
        '2010-12-30'
        >>> v["D5DC2A79C2E4AE92"]["algorithm"]
        1
        >>> v["D5DC2A79C2E4AE92"]["subkeys"][0]["algorithm"]
        1
        """
        keys = {}
        curkey = None

        def parse_pubkey(line, curkey, keys):
            keys[line[4]] = {
                "size": int(line[2]),
                "creation-date": line[5],
                "uids": [],
                "subkeys": [],
                "signatures": [],
                "trust": line[1],
                "algorithm": int(line[3])
            }
            if line[6] != "":
                keys[line[4]]["revocation-date"] = line[5]
            curkey = line[4]
            curkey, keys = parse_uid(line, curkey, keys)
            return (curkey, keys)

        def parse_subkey(line, curkey, keys):
            subkey = {"id": line[4], "size": int(line[2]), "creation-date": line[5], "algorithm": int(line[3])}
            if line[0] == "ssb":
                subkey["secret"] = True
            keys[curkey]["subkeys"].append(subkey)
            return (curkey, keys)

        def parse_fingerprint(line, curkey, keys):
            keys[curkey]["fingerprint"] = line[9]
            return (curkey, keys)

        def parse_userattribute(line, curkey, keys):
            # TODO: We are currently ignoring user attributes as not useful.
            #       We may at some point want to use --attribute-fd and read
            #       in user photos and such?
            return (curkey, keys)

        def parse_privkey(line, curkey, keys):
            curkey, keys = parse_pubkey(line, curkey, keys)
            return (curkey, keys)

        def parse_uid(line, curkey, keys):
            matches = re.match("([^\(\)]*)( \((.*)\)){0,1} \<(.*)\>", line[9])
            if matches:
                email = matches.groups(0)[3]
                comment = matches.groups(0)[2] or ""
                name = matches.groups(0)[0]
            else:
                email = line[9]
                name = ""
                comment = ""
            keys[curkey]["uids"].append({"email": email, "name": name, "comment": comment, "creation-date": line[5] })
            return (curkey, keys)

        def parse_trust(line, curkey, keys):
            # TODO: We are currently ignoring commentary from the Trust DB.
            return (curkey, keys)

        def parse_signature(line, curkey, keys):
            sig = {"signer": line[9], "signature-date": line[5], "keyid": line[4], "trust": line[10], "algorithm": line[4]}

            keys[curkey]["signatures"].append(sig)
            return (curkey, keys)

        def parse_unknown(line, curkey, keys):
            print "Unknown line with code '%s'" % line[0]
            return (curkey, keys)

        def parse_none(line, curkey, keys):
            return (curkey, keys)

        disp = {"pub": parse_pubkey,
                "sub": parse_subkey,
                "ssb": parse_subkey,
                "fpr": parse_fingerprint,
                "uat": parse_userattribute,
                "sec": parse_privkey,
                "tru": parse_trust,
                "sig": parse_signature,
                "uid": parse_uid,
                "gpg": parse_none,
               }

        lines = keylist.split("\n")
        for line in lines:
            if line == "":
                continue
            parms = line.split(":")
            r = disp.get(parms[0], parse_unknown)
            curkey, keys = r(parms, curkey, keys)

        return keys

    def emptycallbackmap():
        """Utility function for people who are confused about what callbacks exist."""
        return dict([[x, []] for x in self.needed_fds])

    def run(self, args=[], callbacks={}, output=None, debug=False):
        """
        >>> g = GnuPG()
        >>> g.run(["--list-keys"])[0]
        0
        """
        self.pipes = {}
        args.insert(0, self.gpgbinary)
        args.append("--with-colons")
        args.append("--verbose")
        args.append("--enable-progress-filter")

        for fd in self.fds.keys():
            if fd not in self.needed_fds:
                continue
            self.pipes[fd] = os.pipe()
            if debug:
                print "Opening fd %s, fh %d, mode %s" % (fd,
                    self.pipes[fd][self.fds[fd]], ["r", "w"][self.fds[fd]])
            args.append("--%s-fd" % fd)
            # The remote end of the pipe:
            args.append("%d" % self.pipes[fd][not self.fds[fd]])
            fdno = self.pipes[fd][self.fds[fd]]
            self.handles[fd] = os.fdopen(fdno, ["r", "w"][self.fds[fd]])
            # Cause file handles to stay open after execing
            fcntl.fcntl(self.handles[fd], fcntl.F_SETFD, 0)
            fl = fcntl.fcntl(self.handles[fd], fcntl.F_GETFL)
            fcntl.fcntl(self.handles[fd], fcntl.F_SETFL, fl | os.O_NONBLOCK)

        if debug: print "Running gpg with %s" % args

        proc = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE)

        self.handles["stdout"] = proc.stdout
        self.handles["stderr"] = proc.stderr
        self.handles["stdin"] = proc.stdin

        if output:
            self.handles["stdin"].write(output)
            self.handles["stdin"].close()

        if self.passphrase:
            self.handles["passphrase"].write(self.passphrase)
            self.handles["passphrase"].close()

        retvals = {}
        for fd in self.needed_fds:
            if fd in ("stdin", "passphrase"):
                continue
            if not callbacks.has_key(fd):
                continue
            while True:
                if debug: print "Reading %s" % fd

                try:
                    buf = self.handles[fd].read()
                except IOError:
                    break

                if not retvals.has_key(fd):
                    retvals[fd] = []
                if buf == "":
                    break

                if type(callbacks[fd]) == list:
                    for cb in callbacks[fd]:
                        retvals[fd].append(cb(buf))
                else:
                    retvals[fd].append(callbacks[fd](buf))

        return proc.returncode, retvals

    def list_keys(self):
        """
        >>> g = GnuPG()
        >>> g.list_keys()[0]
        0
        """
        retvals = self.run(["--list-keys", "--fingerprint"], callbacks={"stdout": self.parse_keylist})
        return retvals

    def list_sigs(self):
        self.run(["--list-sigs", "--fingerprint"], callbacks={"stdout": self.parse_keylist})
        return retvals

    def list_secret_keys(self):
        """
        >>> g = GnuPG()
        >>> g.list_secret_keys()[0]
        0
        """
        retvals = self.run(["--list-secret-keys", "--fingerprint"], callbacks={"stdout": self.parse_keylist})
        return retvals

    def encrypt(self, data, to=[], armor=True):
        """
        >>> g = GnuPG()
        >>> g.encrypt("Hello, World", to=["smari@mailpile.is"])[0]
        0
        """
        action = ["--encrypt"]
        if armor:
            action.append("--armor")
        for r in to:
            action.append("--recipient")
            action.append(r)
        retvals = self.run(action, callbacks={"stdout": self.default_output}, output=data)
        return retvals[0], retvals[1]["stdout"][0]

    def decrypt(self, data, passphrase=None):
        """
        Note that this test will fail if you don't replace the recipient with one whose key you control.
        >>> g = GnuPG()
        >>> ct = g.encrypt("Hello, World", to=["smari@mailpile.is"])[1]
        >>> g.decrypt(ct)[1]
        'Hello, World'
        """
        if passphrase:
            self.passphrase = passphrase
        action = ["--decrypt"]
        retvals = self.run(action, callbacks={"stdout": self.default_output}, output=data)
        return retvals[0], retvals[1]["stdout"][0]

    def sign(self, data, _from=None, armor=True, detatch=True, clearsign=False, passphrase=None):
        """
        >>> g = GnuPG()
        >>> g.sign("Hello, World", _from="smari@mailpile.is")[0]
        0
        """
        if passphrase:
            self.passphrase = passphrase
        if detatch and not clearsign:
            action = ["--detach-sign"]
        elif clearsign:
            action = ["--clearsign"]
        else:
            action = ["--sign"]
        if armor:
            action.append("--armor")
        if _from:
            action.append("--local-user")
            action.append(_from)

        retvals = self.run(action, callbacks={"stdout": self.default_output}, output=data)
        return retvals[0], retvals[1]["stdout"][0]

    def verify(self, data, signature=None):
        """
        >>> g = GnuPG()
        >>> s = g.sign("Hello, World", _from="smari@mailpile.is", clearsign=True)[1]
        >>> g.verify(s)
        """
        print "BEFORE: %s" % self.needed_fds
        self.needed_fds.append("status")
        print "AFTER: %s" % self.needed_fds
        retvals = self.run(["--verify"], callbacks={"stderr": self.parse_verify, "status": self.default_output}, output=data, debug=True)
        self.needed_fds.remove("status")

        return retvals[0], retvals[1]["stderr"]

    def sign_encrypt(self, data, _from=None, to=[], armor=True, detatch=True, clearsign=False):
        retval, signblock = self.sign(data, _from=_from, armor=armor, detatch=detatch, clearsign=clearsign)
        if detatch:
            # TODO: Deal with detached signature.
            retval, cryptblock = self.encrypt(data, to=to, armor=armor)
        else:
            retval, cryptblock = self.encrypt(signblock, to=to, armor=armor)

        return cryptblock

if __name__ == "__main__":
    g = GnuPG()
    g.verify("Foo")
#    import doctest
#    t = doctest.testmod()
#    if t.failed == 0:
#        print "GPG Interface: All %d tests successful" % (t.attempted)
#    else:
#        print "GPG Interface: %d out of %d tests failed" % (t.failed, t.attempted)
<|endoftext|>


<commit_before># Copyright (c) OpenMMLab. All rights reserved.
import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist
from mmcv.utils import get_git_hash

from mmdet import __version__
from mmdet.apis import init_random_seed, set_random_seed, train_detector
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.utils import (collect_env, get_device, get_root_logger,
                         replace_cfg_vals, setup_multi_processes,
                         update_data_root)

def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--auto-resume',
        action='store_true',
        help='resume from the latest checkpoint automatically')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='(Deprecated, please use --gpu-id) number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--diff-seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file (deprecate), '
        'change to --cfg-options instead.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument(
        '--auto-scale-lr',
        action='store_true',
        help='enable automatically scaling LR.')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    if args.options and args.cfg_options:
        raise ValueError(
            '--options and --cfg-options cannot be both '
            'specified, --options is deprecated in favor of --cfg-options')
    if args.options:
        warnings.warn('--options is deprecated in favor of --cfg-options')
        args.cfg_options = args.options

    return args

def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)

    # replace the ${key} with the value of cfg.key
    cfg = replace_cfg_vals(cfg)

    # update data root according to MMDET_DATASETS
    update_data_root(cfg)

    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    if args.auto_scale_lr:
        if 'auto_scale_lr' in cfg and \
                'enable' in cfg.auto_scale_lr and \
                'base_batch_size' in cfg.auto_scale_lr:
            cfg.auto_scale_lr.enable = True
        else:
            warnings.warn('Can not find "auto_scale_lr" or '
                          '"auto_scale_lr.enable" or '
                          '"auto_scale_lr.base_batch_size" in your'
                          ' configuration file. Please update all the '
                          'configuration files to mmdet >= 2.24.1.')

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])

    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    cfg.auto_resume = args.auto_resume
    if args.gpus is not None:
        cfg.gpu_ids = range(1)
        warnings.warn('`--gpus` is deprecated because we only support '
                      'single GPU mode in non-distributed training. '
                      'Use `gpus=1` now.')
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed training. Use the first GPU '
                      'in `gpu_ids` now.')
    if args.gpus is None and args.gpu_ids is None:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        # re-set gpu_ids with distributed training mode
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info
    meta['config'] = cfg.pretty_text
    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    cfg.device = get_device()
    # set random seeds
    seed = init_random_seed(args.seed, device=cfg.device)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = seed
    meta['exp_name'] = osp.basename(args.config)

    model = build_detector(
        cfg.model,
        train_cfg=cfg.get('train_cfg'),
        test_cfg=cfg.get('test_cfg'))
    model.init_weights()

    datasets = [build_dataset(cfg.data.train)]
    if len(cfg.workflow) == 2:
        val_dataset = copy.deepcopy(cfg.data.val)
        val_dataset.pipeline = cfg.data.train.pipeline
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmdet version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmdet_version=__version__ + get_git_hash()[:7],
            CLASSES=datasets[0].CLASSES)
    # add an attribute for visualization convenience
    model.CLASSES = datasets[0].CLASSES
    train_detector(
        model,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        meta=meta)

if __name__ == '__main__':
    main()
<commit_msg>[Fix]: Fix get train_pipeline method of val workflow (#8575)<commit_after># Copyright (c) OpenMMLab. All rights reserved.
import argparse
import copy
import os
import os.path as osp
import time
import warnings

import mmcv
import torch
import torch.distributed as dist
from mmcv import Config, DictAction
from mmcv.runner import get_dist_info, init_dist
from mmcv.utils import get_git_hash

from mmdet import __version__
from mmdet.apis import init_random_seed, set_random_seed, train_detector
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.utils import (collect_env, get_device, get_root_logger,
                         replace_cfg_vals, setup_multi_processes,
                         update_data_root)

def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='train config file path')
    parser.add_argument('--work-dir', help='the dir to save logs and models')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--auto-resume',
        action='store_true',
        help='resume from the latest checkpoint automatically')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='(Deprecated, please use --gpu-id) number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--diff-seed',
        action='store_true',
        help='Whether or not set different seeds for different ranks')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file (deprecate), '
        'change to --cfg-options instead.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument(
        '--auto-scale-lr',
        action='store_true',
        help='enable automatically scaling LR.')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    if args.options and args.cfg_options:
        raise ValueError(
            '--options and --cfg-options cannot be both '
            'specified, --options is deprecated in favor of --cfg-options')
    if args.options:
        warnings.warn('--options is deprecated in favor of --cfg-options')
        args.cfg_options = args.options

    return args

def main():
    args = parse_args()

    cfg = Config.fromfile(args.config)

    # replace the ${key} with the value of cfg.key
    cfg = replace_cfg_vals(cfg)

    # update data root according to MMDET_DATASETS
    update_data_root(cfg)

    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    if args.auto_scale_lr:
        if 'auto_scale_lr' in cfg and \
                'enable' in cfg.auto_scale_lr and \
                'base_batch_size' in cfg.auto_scale_lr:
            cfg.auto_scale_lr.enable = True
        else:
            warnings.warn('Can not find "auto_scale_lr" or '
                          '"auto_scale_lr.enable" or '
                          '"auto_scale_lr.base_batch_size" in your'
                          ' configuration file. Please update all the '
                          'configuration files to mmdet >= 2.24.1.')

    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # work_dir is determined in this priority: CLI > segment in file > filename
    if args.work_dir is not None:
        # update configs according to CLI args if args.work_dir is not None
        cfg.work_dir = args.work_dir
    elif cfg.get('work_dir', None) is None:
        # use config filename as default work_dir if cfg.work_dir is None
        cfg.work_dir = osp.join('./work_dirs',
                                osp.splitext(osp.basename(args.config))[0])

    if args.resume_from is not None:
        cfg.resume_from = args.resume_from
    cfg.auto_resume = args.auto_resume
    if args.gpus is not None:
        cfg.gpu_ids = range(1)
        warnings.warn('`--gpus` is deprecated because we only support '
                      'single GPU mode in non-distributed training. '
                      'Use `gpus=1` now.')
    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed training. Use the first GPU '
                      'in `gpu_ids` now.')
    if args.gpus is None and args.gpu_ids is None:
        cfg.gpu_ids = [args.gpu_id]

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
        # re-set gpu_ids with distributed training mode
        _, world_size = get_dist_info()
        cfg.gpu_ids = range(world_size)

    # create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    # dump config
    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
    # init the logger before other steps
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    # init the meta dict to record some important information such as
    # environment info and seed, which will be logged
    meta = dict()
    # log env info
    env_info_dict = collect_env()
    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
    dash_line = '-' * 60 + '\n'
    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                dash_line)
    meta['env_info'] = env_info
    meta['config'] = cfg.pretty_text
    # log some basic info
    logger.info(f'Distributed training: {distributed}')
    logger.info(f'Config:\n{cfg.pretty_text}')

    cfg.device = get_device()
    # set random seeds
    seed = init_random_seed(args.seed, device=cfg.device)
    seed = seed + dist.get_rank() if args.diff_seed else seed
    logger.info(f'Set random seed to {seed}, '
                f'deterministic: {args.deterministic}')
    set_random_seed(seed, deterministic=args.deterministic)
    cfg.seed = seed
    meta['seed'] = seed
    meta['exp_name'] = osp.basename(args.config)

    model = build_detector(
        cfg.model,
        train_cfg=cfg.get('train_cfg'),
        test_cfg=cfg.get('test_cfg'))
    model.init_weights()

    datasets = [build_dataset(cfg.data.train)]
    if len(cfg.workflow) == 2:
        assert 'val' in [mode for (mode, _) in cfg.workflow]
        val_dataset = copy.deepcopy(cfg.data.val)
        val_dataset.pipeline = cfg.data.train.get(
            'pipeline', cfg.data.train.dataset.get('pipeline'))
        datasets.append(build_dataset(val_dataset))
    if cfg.checkpoint_config is not None:
        # save mmdet version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmdet_version=__version__ + get_git_hash()[:7],
            CLASSES=datasets[0].CLASSES)
    # add an attribute for visualization convenience
    model.CLASSES = datasets[0].CLASSES
    train_detector(
        model,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        meta=meta)

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