<commit_before>"""
   Copyright (C) 2012 STFC

   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 __future__ import print_function

from ssm import crypto
from ssm.message_directory import MessageDirectory

try:
    from dirq.QueueSimple import QueueSimple
    from dirq.queue import Queue
except ImportError:
    # ImportError is raised later on if dirq is requested but not installed.
    QueueSimple = None
    Queue = None

import stomp
from stomp.exception import ConnectFailedException

import os
import socket
import time
from logging import getLogger, INFO, WARNING, DEBUG

try:
    from argo_ams_library import ArgoMessagingService, AmsMessage
except ImportError:
    # ImportError is raised later on if AMS is requested but lib not installed.
    ArgoMessagingService = None
    AmsMessage = None

# Set up logging
log = getLogger(__name__)


class Ssm2Exception(Exception):
    """Exception for use by SSM2."""

    pass


class Ssm2(stomp.ConnectionListener):
    """Minimal SSM implementation."""

    # Schema for the dirq message queue.
    QSCHEMA = {'body': 'string', 'signer': 'string', 'empaid': 'string?'}
    REJECT_SCHEMA = {'body': 'string', 'signer': 'string?',
                     'empaid': 'string?', 'error': 'string'}
    CONNECTION_TIMEOUT = 10

    # Messaging protocols
    STOMP_MESSAGING = 'STOMP'
    AMS_MESSAGING = 'AMS'

    def __init__(self, hosts_and_ports, qpath, cert, key, dest=None, listen=None,
                 capath=None, check_crls=False, use_ssl=False, username=None, password=<PASSWORD>,
                 enc_cert=None, verify_enc_cert=True, pidfile=None, path_type='dirq',
                 protocol=STOMP_MESSAGING, project=None, token=''):
        """Create an SSM2 object.

        If a listen value is supplied, this SSM2 will be a receiver.
        """
        self._conn = None
        self._last_msg = None

        self._brokers = hosts_and_ports
        self._cert = cert
        self._key = key
        self._enc_cert = enc_cert
        self._capath = capath
        self._check_crls = check_crls
        self._user = username
        self._pwd = password
        self._use_ssl = use_ssl
        # use pwd auth if we're supplied both user and pwd
        self._use_pwd = username is not None and password is not None
        self.connected = False

        self._listen = listen
        self._dest = dest

        self._valid_dns = []
        self._pidfile = pidfile

        # Used to differentiate between STOMP and AMS methods
        self._protocol = protocol

        # Used when interacting with an Argo Messaging Service
        self._project = project
        self._token = token

        if self._protocol == Ssm2.AMS_MESSAGING:
            if ArgoMessagingService is None:
                raise ImportError(
                    "The Python package argo_ams_library must be installed to "
                    "use AMS. Please install or use STOMP."
                )
            self._ams = ArgoMessagingService(endpoint=self._brokers[0],
                                             token=self._token,
                                             cert=self._cert,
                                             key=self._key,
                                             project=self._project)

        # create the filesystem queues for accepted and rejected messages
        if dest is not None and listen is None:
            # Determine what sort of outgoing structure to make
            if path_type == 'dirq':
                if QueueSimple is None:
                    raise ImportError("dirq path_type requested but the dirq "
                                      "module wasn't found.")

                self._outq = QueueSimple(qpath)

            elif path_type == 'directory':
                self._outq = MessageDirectory(qpath)
            else:
                raise Ssm2Exception('Unsupported path_type variable.')

        elif listen is not None:
            inqpath = os.path.join(qpath, 'incoming')
            rejectqpath = os.path.join(qpath, 'reject')

            # Receivers must use the dirq module, so make a quick sanity check
            # that dirq is installed.
            if Queue is None:
                raise ImportError("Receiving SSMs must use dirq, but the dirq "
                                  "module wasn't found.")

            self._inq = Queue(inqpath, schema=Ssm2.QSCHEMA)
            self._rejectq = Queue(rejectqpath, schema=Ssm2.REJECT_SCHEMA)
        else:
            raise Ssm2Exception('SSM must be either producer or consumer.')
        # check that the cert and key match
        if not crypto.check_cert_key(self._cert, self._key):
            raise Ssm2Exception('Cert and key don\'t match.')

        # Check that the certificate has not expired.
        if not crypto.verify_cert_date(self._cert):
            raise Ssm2Exception('Certificate %s has expired or will expire '
                                'within a day.' % self._cert)

        # check the server certificate provided
        if enc_cert is not None:
            log.info('Messages will be encrypted using %s', enc_cert)
            if not os.path.isfile(self._enc_cert):
                raise Ssm2Exception('Specified certificate file does not exist: %s.' % self._enc_cert)
            # Check that the encyption certificate has not expired.
            if not crypto.verify_cert_date(enc_cert):
                raise Ssm2Exception(
                    'Encryption certificate %s has expired or will expire '
                    'within a day. Please obtain the new one from the final '
                    'server receiving your messages.' % enc_cert
                )
            if verify_enc_cert:
                if not crypto.verify_cert_path(self._enc_cert, self._capath, self._check_crls):
                    raise Ssm2Exception('Failed to verify server certificate %s against CA path %s.'
                                        % (self._enc_cert, self._capath))

        # If the overall SSM log level is info, we want to only
        # see entries from stomp.py and connectionpool at WARNING and above.
        if getLogger("ssm.ssm2").getEffectiveLevel() == INFO:
            getLogger("stomp.py").setLevel(WARNING)
            getLogger("requests.packages.urllib3.connectionpool"
                      ).setLevel(WARNING)
            getLogger("urllib3.connectionpool").setLevel(WARNING)
        # If the overall SSM log level is debug, we want to only
        # see entries from stomp.py and connectionpool at INFO above.
        elif getLogger("ssm.ssm2").getEffectiveLevel() == DEBUG:
            getLogger("stomp.py").setLevel(INFO)
            getLogger("requests.packages.urllib3.connectionpool"
                      ).setLevel(INFO)
            getLogger("urllib3.connectionpool").setLevel(INFO)

    def set_dns(self, dn_list):
        """Set the list of DNs which are allowed to sign incoming messages."""
        self._valid_dns = dn_list

    ##########################################################################
    # Methods called by stomppy
    ##########################################################################

    def on_send(self, frame, unused_body=None):
        """Log sending of message with empaid if available.

        Called by stomppy when a message is sent. 'unused_body' is only present
        to have a backward compatible method signature for stomp.py v3.1.X
        """
        try:
            # Try the stomp.py v4 way first
            empaid = frame.headers['empa-id']
        except KeyError:
            # Then you are most likely using stomp.py v4.<commit_msg>Factor out logging tweaks to own method<commit_after>"""
   Copyright (C) 2012 STFC

   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 __future__ import print_function

from ssm import crypto
from ssm.message_directory import MessageDirectory

try:
    from dirq.QueueSimple import QueueSimple
    from dirq.queue import Queue
except ImportError:
    # ImportError is raised later on if dirq is requested but not installed.
    QueueSimple = None
    Queue = None

import stomp
from stomp.exception import ConnectFailedException

import os
import socket
import time
from logging import getLogger, INFO, WARNING, DEBUG

try:
    from argo_ams_library import ArgoMessagingService, AmsMessage
except ImportError:
    # ImportError is raised later on if AMS is requested but lib not installed.
    ArgoMessagingService = None
    AmsMessage = None

# Set up logging
log = getLogger(__name__)


class Ssm2Exception(Exception):
    """Exception for use by SSM2."""

    pass


class Ssm2(stomp.ConnectionListener):
    """Minimal SSM implementation."""

    # Schema for the dirq message queue.
    QSCHEMA = {'body': 'string', 'signer': 'string', 'empaid': 'string?'}
    REJECT_SCHEMA = {'body': 'string', 'signer': 'string?',
                     'empaid': 'string?', 'error': 'string'}
    CONNECTION_TIMEOUT = 10

    # Messaging protocols
    STOMP_MESSAGING = 'STOMP'
    AMS_MESSAGING = 'AMS'

    def __init__(self, hosts_and_ports, qpath, cert, key, dest=None, listen=None,
                 capath=None, check_crls=False, use_ssl=False, username=None, password=<PASSWORD>,
                 enc_cert=None, verify_enc_cert=True, pidfile=None, path_type='dirq',
                 protocol=STOMP_MESSAGING, project=None, token=''):
        """Create an SSM2 object.

        If a listen value is supplied, this SSM2 will be a receiver.
        """
        self._conn = None
        self._last_msg = None

        self._brokers = hosts_and_ports
        self._cert = cert
        self._key = key
        self._enc_cert = enc_cert
        self._capath = capath
        self._check_crls = check_crls
        self._user = username
        self._pwd = password
        self._use_ssl = use_ssl
        # use pwd auth if we're supplied both user and pwd
        self._use_pwd = username is not None and password is not None
        self.connected = False

        self._listen = listen
        self._dest = dest

        self._valid_dns = []
        self._pidfile = pidfile

        # Used to differentiate between STOMP and AMS methods
        self._protocol = protocol

        # Used when interacting with an Argo Messaging Service
        self._project = project
        self._token = token

        if self._protocol == Ssm2.AMS_MESSAGING:
            if ArgoMessagingService is None:
                raise ImportError(
                    "The Python package argo_ams_library must be installed to "
                    "use AMS. Please install or use STOMP."
                )
            self._ams = ArgoMessagingService(endpoint=self._brokers[0],
                                             token=self._token,
                                             cert=self._cert,
                                             key=self._key,
                                             project=self._project)

        # create the filesystem queues for accepted and rejected messages
        if dest is not None and listen is None:
            # Determine what sort of outgoing structure to make
            if path_type == 'dirq':
                if QueueSimple is None:
                    raise ImportError("dirq path_type requested but the dirq "
                                      "module wasn't found.")

                self._outq = QueueSimple(qpath)

            elif path_type == 'directory':
                self._outq = MessageDirectory(qpath)
            else:
                raise Ssm2Exception('Unsupported path_type variable.')

        elif listen is not None:
            inqpath = os.path.join(qpath, 'incoming')
            rejectqpath = os.path.join(qpath, 'reject')

            # Receivers must use the dirq module, so make a quick sanity check
            # that dirq is installed.
            if Queue is None:
                raise ImportError("Receiving SSMs must use dirq, but the dirq "
                                  "module wasn't found.")

            self._inq = Queue(inqpath, schema=Ssm2.QSCHEMA)
            self._rejectq = Queue(rejectqpath, schema=Ssm2.REJECT_SCHEMA)
        else:
            raise Ssm2Exception('SSM must be either producer or consumer.')
        # check that the cert and key match
        if not crypto.check_cert_key(self._cert, self._key):
            raise Ssm2Exception('Cert and key don\'t match.')

        # Check that the certificate has not expired.
        if not crypto.verify_cert_date(self._cert):
            raise Ssm2Exception('Certificate %s has expired or will expire '
                                'within a day.' % self._cert)

        # check the server certificate provided
        if enc_cert is not None:
            log.info('Messages will be encrypted using %s', enc_cert)
            if not os.path.isfile(self._enc_cert):
                raise Ssm2Exception('Specified certificate file does not exist: %s.' % self._enc_cert)
            # Check that the encyption certificate has not expired.
            if not crypto.verify_cert_date(enc_cert):
                raise Ssm2Exception(
                    'Encryption certificate %s has expired or will expire '
                    'within a day. Please obtain the new one from the final '
                    'server receiving your messages.' % enc_cert
                )
            if verify_enc_cert:
                if not crypto.verify_cert_path(self._enc_cert, self._capath, self._check_crls):
                    raise Ssm2Exception('Failed to verify server certificate %s against CA path %s.'
                                        % (self._enc_cert, self._capath))

        self._set_external_logging_levels()

    def _set_external_logging_levels(self):
        """Tweak the logging of dependencies to better match SSM verbosity."""
        # If the overall SSM log level is INFO, we want to only
        # see entries from stomp.py and connectionpool at WARNING and above.
        if getLogger("ssm.ssm2").getEffectiveLevel() == INFO:
            getLogger("stomp.py").setLevel(WARNING)
            getLogger("requests.packages.urllib3.connectionpool"
                      ).setLevel(WARNING)
            getLogger("urllib3.connectionpool").setLevel(WARNING)
        # If the overall SSM log level is DEBUG, we want to only
        # see entries from stomp.py and connectionpool at INFO above.
        elif getLogger("ssm.ssm2").getEffectiveLevel() == DEBUG:
            getLogger("stomp.py").setLevel(INFO)
            getLogger("requests.packages.urllib3.connectionpool"
                      ).setLevel(INFO)
            getLogger("urllib3.connectionpool").setLevel(INFO)

    def set_dns(self, dn_list):
        """Set the list of DNs which are allowed to sign incoming messages."""
        self._valid_dns = dn_list

    ##########################################################################
    # Methods called by stomppy
    ##########################################################################

    def on_send(self, frame, unused_body=None):
        """Log sending of message with empaid if available.

        Called by stomppy when a message is sent. 'unused_body' is only present
        to have a backward compatible method signature for stomp.py v3.1.X
        """
        try:
            # Try the stomp.py v4 way first
            empaid = frame.headers['empa-id']
        except KeyError:
            # Then you are most likely using stomp.py v4.
{'fraction_non_alphanumeric': 0.08964554503520272, 'fraction_numerical': 0.004369992716678805, 'mean_word_length': 2.1035976643435674, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4062268', 'n_tokens_mistral': 4399, 'n_tokens_neox': 4205, 'n_words': 1462}