PORT = 5872
NONCE_LEN = 16
NULL_NONCE = b'\x00' * NONCE_LEN
class DiamondRPCFailure(RPCError):
class DiamondRPCFCacheMiss(RPCError):
class DiamondRPCCookieExpired(RPCError):
class DiamondRPCSchemeNotSupported(RPCError):
class XDR_attribute(XDRStruct):
class XDR_object(XDRStruct):
class XDR_blob_list(XDRStruct):
class XDR_filter_config(XDRStruct):
class XDR_setup(XDRStruct):
class XDR_blob_data(XDRStruct):
class XDR_start(XDRStruct):
class XDR_stat(XDRStruct):
class XDR_filter_stats(XDRStruct):
class XDR_search_stats(XDRStruct):
class XDR_session_var(XDRStruct):
class XDR_session_vars(XDRStruct):
class XDR_reexecute(XDRStruct):
class XDR_retrain(XDRStruct):
class XDR_attribute_list(XDRStruct): class ExecutableBlobCache(BlobCache):
    '''A BlobCache that can create executable files from cache entries.

    Creates a temporary directory in TMPDIR and does not clean it up,
    under the expectation that the entire TMPDIR will be blown away after
    the search is complete.
    '''

    def __init__(self, basedir):
        BlobCache.__init__(self, basedir)
        # Ensure _executable_dir is inside the search-specific tempdir
        self._executable_dir = mkdtemp(dir=os.environ.get('TMPDIR'),
                                       prefix='executable-')

    def executable_path(self, sig):
        '''Return a path to the file containing the specified data
        so that the file can be executed.  Ensure that the executable
        bit is set in the filesystem.  The path is guaranteed to be valid
        for the lifetime of the search, even in the presence of
        garbage-collection.'''
        src = self._path(sig)
        dest = os.path.join(self._executable_dir, sig)

        def make_dest():
            # Link the blob into a temporary directory.  This directory will
            # normally (but need not always) be deleted by the supervisor when
            # the search terminates.
            try:
                # We have to use mktemp() here because link() requires the
                # destination not to exist.  If someone else wins the race
                # to create the same file we'll simply fall back to
                # copying it.
                dest_tmp = mktemp(dir=self._executable_dir)
                os.link(src, dest_tmp)
            except OSError:
                # self._executable_dir may be on a different filesystem
                # than self.basedir.  Try copying the file instead.
                fd, dest_tmp = mkstemp(dir=self._executable_dir)
                os.close(fd)
                shutil.copyfile(src, dest_tmp)
            os.chmod(dest_tmp, 0o500)
            os.rename(dest_tmp, dest)
        self._access(sig)
        if not os.path.exists(dest):
            self._try_with_rescue(sig, make_dest, (OSError, IOError))
        return dest class DiamondRPCFailure(RPCError):
    '''Generic Diamond RPC failure.'''
    code = 500 class DiamondRPCFCacheMiss(RPCError):
    '''Filter code or blob argument missed in the blob cache.'''
    code = 501 class DiamondRPCCookieExpired(RPCError):
    '''Proffered scope cookie has expired.'''
    code = 504 class DiamondRPCSchemeNotSupported(RPCError):
    '''URI scheme not supported.'''
    code = 505 class RPCHandlers(with_metaclass(_RPCMeta, object)):
    '''Base class of RPC handler objects.'''
    log_rpcs = False

    @staticmethod
    def handler(cmd, request_class=None, reply_class=None):
        '''Decorator declaring the function to be an RPC handler with the
        given command number and request class.'''
        def decorator(func):
            func.rpc_procedure = cmd
            func.rpc_request_class = request_class
            func.rpc_reply_class = reply_class
            return func
        return decorator

    # self._cmds is created by _RPCMeta
    # pylint: disable=no-member
    def get_handler(self, procedure):
        '''Returns the handler function for the specified procedure number
        or raises KeyError.'''
        # self._cmds contains unbound methods; return a bound one
        return self._cmds[procedure].__get__(self, self.__class__)
    # pylint: enable=no-member class RPCError(Exception):
    '''Base class for RPC error codes.'''
    code = None

    # pylint doesn't know about __subclasses__
    # pylint: disable=no-member
    @classmethod
    def get_class(cls, code):
        '''Return the error subclass for the specified code.'''
        for c in cls.__subclasses__():
            if c.code == code:
                return c
            try:
                return c.get_class(code)
            except KeyError:
                pass
        raise KeyError('Error code %d' % code)
    # pylint: enable=no-member class RPCProcedureUnavailable(RPCError):
    '''Remote procedure not available.'''
    code = -3 class ScopeCookie(object):
    def __init__(self, serial, expires, blaster, servers, scopeurls, data,
                 signature):
        '''Do not call this directly; use generate() or parse() instead.'''
        assert isinstance(data, str)
        assert isinstance(signature, bytes)
        # Ensure the expiration time is tz-aware
        if expires.tzinfo is None or expires.tzinfo.utcoffset(expires) is None:
            raise ScopeError('Expiration time does not include time zone')
        self.serial = serial            # A UUID object
        self.expires = expires          # A datetime object
        self.blaster = blaster          # The URL of the JSON blaster or None
        self.servers = servers          # A list
        self.scopeurls = scopeurls      # The list of scope URLs
        self.data = data                # All of the above, as a string
        self.signature = signature      # Binary signature of the data (bytes)

    def __str__(self):
        '''Return the decoded scope cookie.'''
        return binascii.hexlify(self.signature).decode() + '\n' + self.data

    def __repr__(self):
        return ('<ScopeCookie %s, blaster %s, servers %s, urls %s, expiration %s>' %
                (self.serial, self.blaster, self.servers, self.scopeurls, self.expires))

    def __iter__(self):
        '''Return an iterator over the scope URLs.'''
        return iter(self.scopeurls)

    def encode(self):
        '''Return the encoded scope cookie as a string.'''
        return (BOUNDARY_START +
                textwrap.fill(base64.b64encode(str(self).encode()).decode(), 64) + '\n' +
                BOUNDARY_END)

    def verify(self, servernames, certdata):
        '''Verify the cookie against the specified server ID list and
        certificate data.'''
        assert isinstance(certdata, str)

        now = datetime.now(tzutc())
        # Check cookie expiration
        if self.expires < now:
            raise ScopeCookieExpired('Cookie expired on %s' % self.expires)
        # Check that one of our server names matches a name in the cookie
        if not set(servernames) & set(self.servers):
            raise ScopeError('Cookie does not contain matching server name')
        # Split certdata into individual certificates
        begin = '-----BEGIN CERTIFICATE-----\n'
        end = '-----END CERTIFICATE-----'        # no trailing newline
        certdata = re.findall(begin + BASE64_RE + end, certdata)
        # Load certificates
        certs = [X509.load_cert_string(cd) for cd in certdata]
        # Find a certificate that verifies the signature
        failure = "Couldn't validate scope cookie signature"
        for cert in certs:
            # Check the signature
            key = cert.get_pubkey()
            key.verify_init()
            key.verify_update(self.data.encode())
            if key.verify_final(self.signature) != 1:
                # Signature does not match certificate
                continue
            # Signature valid; now check the certificate
            if cert.verify() != 1:
                failure = 'Scope cookie signed by invalid certificate'
                continue
            if cert.get_not_before().get_datetime() > now:
                failure = 'Scope cookie signed by postdated certificate'
                continue
            if cert.get_not_after().get_datetime() < now:
                failure = 'Scope cookie signed by expired certificate'
                continue
            # We have a match
            return
        raise ScopeError(failure)

    @classmethod
    def generate(cls, servers, scopeurls, expires, keydata, blaster=None):
        '''Generate and return a new ScopeCookie.  servers and scopeurls
        are lists of strings, already Punycoded/URL-encoded as appropriate.
        expires is a timezone-aware datetime.  keydata is a PEM-encoded
        private key.  blaster is an optional string, already URL-encoded.'''
        # Unicode strings can cause signature validation errors
        if isinstance(keydata, str):
            keydata = keydata.encode()
        else:
            assert isinstance(keydata, bytes)

        servers = [str(s) for s in servers]
        scopeurls = [str(u) for u in scopeurls]
        if blaster is not None:
            blaster = str(blaster)
        # Generate scope data
        serial = str(uuid.uuid4())
        headers = [('Version', COOKIE_VERSION),
                   ('Serial', serial),
                   ('Expires', expires.isoformat()),
                   ('Servers', ';'.join(servers))]
        if blaster is not None:
            headers.append(('Blaster', blaster))
        hdrbuf = ''.join('%s: %s\n' % (k, v) for k, v in headers)
        data = hdrbuf + '\n' + '\n'.join(scopeurls) + '\n'
        # Load the signing key
        key = EVP.load_key_string(keydata)  # expects bytes, not str
        # Sign the data
        key.sign_init()
        key.sign_update(data.encode())
        sig = key.sign_final()
        # Return the scope cookie
        return cls(serial, expires, blaster, servers, scopeurls, data, sig)

    @classmethod
    def parse(cls, data):
        """Parse a (single) scope cookie string and return a ScopeCookie
        
        Arguments:
            data {str} -- A single base64-encoded cookie

        Returns:
            [ScopeCookie] -- [description]
        """
        assert isinstance(data, str)

        # Check for boundary markers and remove them
        match = re.match(BOUNDARY_START + '(' + BASE64_RE + ')' +
                         BOUNDARY_END, data)
        if match is None:
            raise ScopeError('Invalid boundary markers')
        data = match.group(1)
        # Base64-decode contents
        try:
            data = base64.b64decode(data).decode()
        except TypeError:
            raise ScopeError('Invalid Base64 data')
        # Split signature, header, and body
        try:
            signature, data = data.split('\n', 1)
            header, body = data.split('\n\n', 1)
        except ValueError:
            raise ScopeError('Malformed cookie')
        # Decode signature
        try:
            signature = binascii.unhexlify(signature)
        except TypeError:
            raise ScopeError('Malformed signature')
        # Parse headers
        blaster = None
        for line in header.splitlines():
            k, v = line.split(':', 1)
            v = v.strip()
            if k == 'Version':
                try:
                    if int(v) != COOKIE_VERSION:
                        raise ScopeError('Unknown cookie version %s' % v)
                except ValueError:
                    raise ScopeError('Invalid cookie version string')
            elif k == 'Serial':
                try:
                    serial = uuid.UUID(v)
                except ValueError:
                    raise ScopeError('Invalid serial format')
            # Ignore KeyId; we have no way to generate one for new cookies
            elif k == 'Expires':
                try:
                    # Parsing ISO8601 is hard; we can't just use strptime.
                    # If the value does not specify a timezone, the
                    # class constructor will error out.
                    expires = dateutil.parser.parse(v)
                except ValueError:
                    raise ScopeError('Invalid date format')
            elif k == 'Servers':
                servers = [s.strip() for s in re.split('[;,]', v)
                           if s.strip() != '']
            elif k == 'Blaster':
                blaster = v
        # Parse body
        scopeurls = [s for s in [u.strip() for u in body.split('\n')]
                     if s != '']
        # Build scope cookie object
        try:
            return cls(serial, expires, blaster, servers, scopeurls, data,
                       signature)
        except NameError:
            raise ScopeError('Missing cookie header')

    @classmethod
    def split(cls, data):
        '''Split the megacookie data into individual encoded scope cookies.'''
        return re.findall(BOUNDARY_START + BASE64_RE + BOUNDARY_END, data) class ScopeError(Exception):
    '''Error generating, parsing, or verifying scope cookie.''' class ScopeCookieExpired(ScopeError):
    '''Scope cookie has expired.''' class FilterStack(object):
    '''A set of filters which collectively decide to accept or drop an
    object.'''

    def __init__(self, filters=None):
        if filters is None:
            filters = []
        # name -> Filter
        self._filters = dict([(f.name, f) for f in filters])
        # Ordered list of filters to execute
        self._order = list()

        # Resolve declared dependencies
        # Filters we have already resolved
        resolved = set()
        # Filters we are currently resolving
        inprocess = set()

        self._optimized = False

        def resolve(filter):
            if filter in resolved:
                return
            if filter in inprocess:
                raise FilterDependencyError('Circular dependency involving '
                                            + filter.name)
            inprocess.add(filter)
            for depname in filter.dependencies:
                try:
                    resolve(self._filters[depname])
                except KeyError:
                    raise FilterDependencyError('No such filter: ' + depname)
            inprocess.remove(filter)
            self._order.append(filter)
            resolved.add(filter)

        for filter in filters:
            resolve(filter)

    def __len__(self):
        return len(self._order)

    def __iter__(self):
        return iter(self._order)

    def bind(self, state, obj_queue, name='Filter'):
        '''Return a FilterStackRunner that can be used to process objects
        with this filter stack.'''
        return FilterStackRunner(state, self._order, obj_queue, name)

    def start_threads(self, state, count, scope):
        '''Start count threads to process objects with this filter stack.'''

        workers = list()
        obj_queue = mp.JoinableQueue(1000)
        obj_queue.cancel_join_thread()  # avoid hanging if workers die first

        def enqueue_scope(q, scope):
            for obj in scope:
                q.put(obj)
            # wait until all objects processed and close the blast channel by sending an empty object
            q.join()
            state.blast.close()

        t = threading.Thread(target=enqueue_scope, args=(obj_queue, scope), name='enqueue-scope-thread')
        t.daemon = True # make sure the whole process terminates when the main thread exits

        for i in range(count):
            w = self.bind(state, obj_queue, 'Filter-%d' % i)
            w.start()
            _log.debug("Started worker %d", w.pid)
            workers.append(w)

        t.start()

        return workers

    def optimize(self):
        """Optimize execution order of filters"""
        # De-duplicate filter based on cache_digest
        # (assuming identical cache_digest implies identical filters)
        if self._optimized:
            return

        new_order = list()
        seen_digests = set()
        for filter in self._order:
            assert isinstance(filter, Filter)
            if not filter.cache_digest:
                raise FilterDependencyError("Should not optimize filter stack "
                                            "before resolving the filters. "
                                            "Report to the developer.")
            if filter.cache_digest not in seen_digests:
                new_order.append(filter)
                seen_digests.add(filter.cache_digest)
            else:
                _log.info("Filter name %s is removed due to de-duplication.", filter.name)

        # commit
        self._order = new_order
        self._optimized = True

        # If you are going to add some smartness of static "filter ordering",
        # work here. class Filter(object):
    '''A filter with arguments.'''

    def __init__(self, name, code_source, blob_source, min_score, max_score,
                 arguments, dependencies):
        self.name = name
        self.code_source = code_source
        self.blob_source = blob_source
        self.min_score = min_score
        self.max_score = max_score
        self.arguments = arguments
        self.dependencies = dependencies
        self.stats = FilterStatistics(name)

        # Will be initialized during resolve()
        self.code_path = None
        self.signature = None
        self.blob = None
        self.cache_digest = None
        self.mode = None

    def connect(self):
        """Return a FilterConnection. To be reloaded during resolve()"""
        assert self is not None
        raise FilterDependencyError('Filter mode is not resolved yet.')

    def resolve(self, state):
        '''Ensure filter code and blob argument are available in the blob
        cache, load the blob argument, and initialize the cache digest.'''
        if self.code_path is not None:
            return
        # Get path to filter code
        code_path, code_signature = self._resolve_code(state)
        # Get contents of blob argument
        blob, blob_signature = self._resolve_blob(state)
        # Initialize digest
        summary = ([code_signature, self.name] + self.arguments +
                   [blob_signature])
        cache_digest = murmur(' '.join(summary))
        # Commit
        self.code_path = code_path
        self.signature = code_signature
        self.blob = blob
        self.cache_digest = cache_digest

        # Resolve mode and reload connect()
        connector = self._resolve_mode(state)
        self.connect = connector.__get__(self, None)  # bind to instance

    def _resolve_code(self, state):
        '''Returns (code_path, signature).'''
        scheme, path = split_scheme(self.code_source)
        if scheme == 'sha256':
            sig = path.lower()
            try:
                # FIXME with filter mode introduced, not necessary to set the
                # executable bit so early.
                return (state.blob_cache.executable_path(sig), sig)
            except KeyError:
                raise FilterDependencyError('Missing code for filter ' +
                                            self.name)
        else:
            raise FilterUnsupportedSource()

    def _resolve_blob(self, state):
        '''Returns (blob data, signature).'''
        scheme, path = split_scheme(self.blob_source)
        if scheme == 'sha256':
            sig = path.lower()
            try:
                return (state.blob_cache[sig], sig)
            except KeyError:
                raise FilterDependencyError('Missing blob for filter ' +
                                            self.name)
        else:
            raise FilterUnsupportedSource()

    def _resolve_mode(self, state):
        """ Return an unbound method which returns a FilterConnection. """
        assert self.code_path is not None

        def scan_mode(filter_file):
            """Scan the first 100 bytes of file for special tags."""
            first_line = filter_file.read(100)
            if 'diamond-docker-filter' in first_line:
                return 'docker'
            else:
                return 'default'

        self.mode = scan_mode(open(self.code_path, 'r'))

        _log.info('%s: %s', self.name, self.mode)

        if self.mode == 'default':
            # default executable mode
            # TODO handle debug command
            def wrapper(_):
                return _FilterProcess(
                    code_argv=[self.code_path],
                    name=self.name,
                    args=self.arguments,
                    blob=self.blob
                )
        elif self.mode == 'docker':
            # Docker filter listens on TCP port
            try:
                with open(self.code_path, 'r') as f:
                    config = yaml.full_load(f.read())
                docker_image = config['docker_image']
                filter_command = config['filter_command']
                connect_method = config.get('connect_method', 'default')

            except Exception as e:
                raise FilterDependencyError(e)

            _log.info('%s: docker mode: %s', self.name, connect_method)

            docker_port = 5555
            if connect_method == 'default': # TODO phase out support for socat connection. Use fifo for legacy filters.
                # connect through socat
                docker_command = 'socat TCP4-LISTEN:%d,fork,nodelay ' \
                                    'EXEC:\"%s --filter\"' % (docker_port, filter_command)

                def wrapper(_):
                    uri = state.context.ensure_resource('docker',docker_image, docker_command)

                    sock = None
                    host, port = uri['IPAddress'], docker_port
                    for _ in range(10):
                        try:
                            # OS may give up with its own timeout regardless of timeout here
                            sock = socket.create_connection((host, port), 1.0)
                            break
                        except socket.error:
                            sock = None
                            time.sleep(0.5)
                            continue

                    if sock is None:
                        raise FilterExecutionError('Unable to connect to filter at %s: %d' % (host, port))

                    return _FilterTCP(
                        sock=sock,
                        name=self.name,
                        args=self.arguments,
                        blob=self.blob)

            elif connect_method == 'fifo':  # named pipe, try to share containers

                def wrapper(_):
                    uid = str(uuid.uuid4())
                    map_vol = '/diamond-tmp/'
                    # make a pair of fifo in tmp directory. "in" means into container.
                    TMPDIR =  os.getenv('TMPDIR')
                    fifo_in = 'fifo-' + uid + '-in'
                    fifo_out = 'fifo-' + uid + '-out'
                    os.mkfifo(os.path.join(TMPDIR, fifo_in))
                    os.mkfifo(os.path.join(TMPDIR, fifo_out))
                    _log.debug("Created in pipe: %s", os.path.join(TMPDIR, fifo_in))
                    _log.debug("Created out pipe: %s", os.path.join(TMPDIR, fifo_out))

                    # map volume when creating Docker container
                    # keep the container running alive, detached
                    docker_command = '/bin/bash'
                    handle = state.context.ensure_resource('docker', docker_image, docker_command,
                                                           volumes = {TMPDIR: {'bind': map_vol, 'mode': 'rw'}}, tty=True)

                    # launch a new filter process inside the containers
                    client = docker.from_env()
                    container = client.containers.get(handle['name'])
                    exec_command = 'sh -c \"%s --filter  <%s >%s\"' % (filter_command, os.path.join(map_vol, fifo_in), os.path.join(map_vol, fifo_out))
                    _log.debug('docker-exec in %s: %s', handle['name'], exec_command)
                    container.exec_run(cmd=exec_command, stdout=False, stderr=False, stdin=False, tty=True, detach=True)

                    # note the in/out are flipped here.
                    # For unknown reasons I must use os.O_RDWR here to avoid blocking
                    fout = os.fdopen(os.open(os.path.join(TMPDIR, fifo_in), os.O_RDWR), 'wb')
                    fin = os.fdopen(os.open(os.path.join(TMPDIR, fifo_out), os.O_RDWR), 'rb')
                    return _FilterConnection(fin=fin, fout=fout, name=self.name, args = self.arguments, blob=self.blob)

            else:
                raise FilterDependencyError('Unknown connect_method: %s' % connect_method)
        else:
            raise FilterDependencyError('Contact the developer.')

        return wrapper

    @classmethod
    def source_available(cls, state, uri):
        '''Verify the URI to ensure that its data is accessible.  Return
        True if we can access the data, False if we can't and should inform
        the client to that effect.  Raise FilterUnsupportedSource if we don't
        support the URI scheme.'''
        scheme, path = split_scheme(uri)
        if scheme == 'sha256':
            return path in state.blob_cache
        else:
            raise FilterUnsupportedSource()

    def bind(self, state):
        '''Return a _FilterRunner for this filter.'''
        # resolve() must be called first
        assert self.code_path is not None
        return _FilterRunner(state, self) class FilterDependencyError(Exception):
    '''Error processing filter dependencies.''' class FilterUnsupportedSource(Exception):
    '''URI scheme for code or blob source is not supported.''' class EmptyObject(object):
    '''An immutable Diamond object with no data and no attributes.'''

    def __init__(self):
        self._attrs = dict()
        self._signatures = dict()
        self._omit_attrs = set()

    def __str__(self):
        return ''

    def __repr__(self):
        return '<EmptyObject>'

    def __iter__(self):
        '''Return an iterator over the attribute names.'''
        return iter(self._attrs.keys())

    def __contains__(self, key):
        return key in self._attrs

    def __getitem__(self, key):
        return self._attrs[key]

    def __setitem__(self, key, value):
        raise TypeError()

    def get_signature(self, key):
        '''Return the MD5 hash of the attribute value.'''
        return self._signatures[key]

    def omit(self, key):
        '''Record that the attribute is not to be returned to the client.'''
        if key in self:
            self._omit_attrs.add(key)
        else:
            raise KeyError()

    def xdr_attributes(self, output_set=None, for_drop=False):
        '''Return a list of XDR_attribute.'''
        if for_drop:
            # Reexecution dropped the object.  Only encode the object ID.
            send_keys = set([ATTR_OBJ_ID])
        else:
            # Don't encode any evidence of omit attributes.
            send_keys = set(self._attrs.keys()) - self._omit_attrs
        # If we have an output set, only send values for send_keys that are
        # in it.  Otherwise, send values for all send_keys.
        if output_set is not None:
            # Ensure the object ID is always sent.  Don't modify the output
            # set in place.
            output_set = output_set.union([ATTR_OBJ_ID])
            send_values = send_keys.intersection(output_set)
        else:
            send_values = send_keys
        _log.debug("{}: Sending attrs: {}".format(str(self), ','.join(send_values)))
        # Serialize
        attrs = []
        for name in send_keys:
            if name in send_values:
                value = self._attrs[name]
                if not isinstance(value, bytes):
                    value = str(value).encode() # to bytes
            else:
                value = b''
            attrs.append(XDR_attribute(name, value))
        return attrs

    def xdr(self, output_set=None):
        '''Return an XDR_object.'''
        return XDR_object(self.xdr_attributes(output_set))

    def debug(self):
        """Print all attributes"""
        print(str(self))
        print('<internal> %s: %s' % ('src', getattr(self, 'src', None)))
        print('<internal> %s: %s' % ('meta', getattr(self, 'meta', None)))
        for k, v in self._attrs.items():
            if k:   # Skip data attribute
                print('%s: %s' % (k, v)) class Object(EmptyObject):
    '''A mutable Diamond object.'''

    def __init__(self, server_id, url, src=None, meta=None, compute_signature=True):
        EmptyObject.__init__(self)
        self._id = url
        self._compute_signature = compute_signature

        if src:
            self.src = src
        if meta:
            self.meta = meta

        # Set default attributes
        self[ATTR_DEVICE_NAME] = server_id + '\0'
        self[ATTR_OBJ_ID] = url + '\0'

    def __str__(self):
        '''Return the object ID.'''
        return self._id

    def __repr__(self):
        return '<Object %s>' % self

    def __setitem__(self, key, value):
        self._attrs[key] = value
        if self._compute_signature:
            self._signatures[key] = murmur(value)
        else:
            self._signatures[key] = 0 class ObjectLoader(object):
    """A context for populating an Object from the dataretriever.  Allows
    network connections to be reused to fetch multiple objects.  Must not
    be used by more than one thread.

    Attributes are also loaded. There are several possible places attributes
    are loaded from. The precedence order is (later overwrites earlier):
    - 'meta' URI from scope list <object /> element, which can either be a URL or
    local file path.
    - HTTP headers starting with ATTR_HEADER_PREFIX='x-attr-', when retrieving
    the object itself. (happen only when retrieving from data retriever)
    - HTTP header ATTR_HEADER_URL='x-attributes', when retrieving the object itself.
    (happen only when retrieving from data retriever)
    """

    def __init__(self, config, blob_cache):
        self._http = _HttpLoader(config)
        self._blob_cache = blob_cache

    def source_available(self, obj):
        '''Examine the Object and return whether we think we will be able
        to load it, without actually doing so.  This can be used during
        reexecution to determine whether we should return
        DiamondRPCFCacheMiss to the client.'''
        scheme, path = split_scheme(str(obj))
        if scheme == 'sha256':
            return path in self._blob_cache
        # Assume we can always load other types of URLs
        return True

    def load(self, obj):
        '''Retrieve the Object and update it with the information we
        receive.'''

        # 0. Load src URI if not exist (mainly used for re-execution)
        self._load_src_maybe_meta(obj)

        # 1. Load object metadata from meta URI
        if hasattr(obj, 'meta'):
            uri = str(obj.meta)
            scheme, path = split_scheme(uri)
            if scheme == 'file':
                self._load_attrbutes_localfile(obj, path)
            else:  # assume http
                self._load_attributes(obj, uri)

        # 2. Load object content from src URI
        uri = str(obj.src)
        scheme, path = split_scheme(uri)
        if scheme == 'sha256':
            self._load_blobcache(obj, path)
        elif scheme == 'file':
            self._load_localfile(obj, path)
        else:  # assume http
            self._load_dataretriever(obj, uri)

        # Set display name if not already in initial attributes
        if ATTR_DISPLAY_NAME not in obj:
            obj[ATTR_DISPLAY_NAME] = str(obj) + '\0'

    def _load_blobcache(self, obj, signature):
        # Load the object data
        try:
            obj[ATTR_DATA] = self._blob_cache[signature]
        except KeyError:
            raise ObjectLoadError('Object not in cache')

    def _load_dataretriever(self, obj, url):
        headers, body = self._http.get(url)
        # Load the object data
        obj[ATTR_DATA] = body
        # Process loose initial attributes
        for key, value in headers.items():
            if key.lower().startswith(ATTR_HEADER_PREFIX):
                key = key[ATTR_HEADER_PREFIX_LEN:]
                obj[key] = value + '\0'
        # Fetch additional initial attributes if specified
        if ATTR_HEADER_URL in headers:
            attr_url = urljoin(url, headers[ATTR_HEADER_URL])
            self._load_attributes(obj, attr_url)

    def _load_localfile(self, obj, path):
        with open(path, 'rb') as f:
            obj[ATTR_DATA] = f.read()

    # The return type of json.loads() confuses pylint
    # pylint: disable=maybe-no-member
    def _load_attributes(self, obj, url):
        '''Load JSON-encoded attribute data from the specified URL.'''
        _headers, body = self._http.get(url)
        try:
            attrs = json.loads(body)
            if not isinstance(attrs, dict):
                raise ObjectLoadError("Failed to retrieve object attributes")
        except ValueError as e:
            raise ObjectLoadError(str(e))
        for k, v in attrs.items():
            obj[k] = str(v) + '\0'

    # pylint: enable=maybe-no-member

    def _load_attrbutes_localfile(self, obj, path):
        try:
            with open(path, 'r') as f:
                attrs = json.load(f)
                if not isinstance(attrs, dict):
                    raise ObjectLoadError(
                        "Failed to retrieve object attributes")
        except IOError as e:
            raise ObjectLoadError(str(e))
        except ValueError as e:
            raise ObjectLoadError(str(e))
        for k, v in attrs.items():
            obj[k] = str(v) + '\0'

    def _load_src_maybe_meta(self, obj):
        """
        Load the src and optional meta URIs from the object's id.
        :param obj:
        :return:
        """
        if not hasattr(obj, 'src'):
            uri = str(obj)
            scheme, path = split_scheme(uri)
            if scheme == 'sha256':
                # Blob has no meta URI
                obj.src = uri
            else:
                # Assume http
                _, body = self._http.get(uri)
                object_el = xmltodict.parse(body)
                obj.src = urljoin(uri, object_el['object']['@src'])
                if '@meta' in object_el['object']:
                    obj.meta = urljoin(uri, object_el['object']['@meta'])
        return class ScopeListLoader(object):
    '''Iterator over the objects in the scope lists referenced by the scope
    cookies.'''

    def __init__(self, config, server_id, cookies):
        self.server_id = server_id
        self.cookies = cookies
        self._config = config
        self._lock = threading.Lock()
        self._handler = _ScopeListHandler()
        self._generator = self._generator_func()

    def __iter__(self):
        return self

    def __next__(self):
        '''Return the next Object.'''
        with self._lock:
            pending_object, scope_url = next(self._generator)

        id = pending_object.pop('id', None)
        src = pending_object.pop('src', None)
        meta = pending_object.pop('meta', None)

        # Allow 'src' and 'meta' be missing
        # If so, they should be loaded later in ObjectLoader        
        if src:
            src = urljoin(scope_url, src)
        if meta:
            meta = urljoin(scope_url, meta)

        # 'src' is the fallback value for 'id' for backward
        # compatibility (f.i. scopelists from Algum)
        # id can't be None eventually
        if id:
            id = urljoin(scope_url, id)
        elif src:
            id = src
        else:
            _log.error('An object cannot have none id and none src at the same time.')

        # if result cache is not enabled, no need to keep signatures in objects
        new_obj = Object(self.server_id,
                            id,
                            src=src,
                            meta=meta,
                            compute_signature=(self._config.cache_server is not None))

        # use the remaining attrs as normal object attributes
        for k, v in pending_object.items():
            new_obj[k] = v + '\0'

        return new_obj

    def _generator_func(self):
        # Build URL opener
        handlers = []
        if self._config.http_proxy is not None:
            handlers.append(urllib.request.ProxyHandler({
                'http': self._config.http_proxy,
                'https': self._config.http_proxy,
            }))
        opener = urllib.request.build_opener(*handlers)
        opener.addheaders = [('User-Agent', self._config.user_agent)]
        # Build XML parser
        parser = make_parser()
        parser.setContentHandler(self._handler)
        # Walk scope URLs
        count = 0
        for scope_url in (url for cookie in self.cookies for url in cookie):
            scope_url = urljoin(BASE_URL, scope_url)
            try:
                # We use urllib2 here because different parts of a single
                # HTTP response will be handled from different threads.
                # pycurl does not support this.
                fh = opener.open(scope_url)
                # Read the scope list in 4 KB chunks
                while True:
                    buf = fh.read(4096)
                    if not buf:
                        break
                    parser.feed(buf)
                    while self._handler.pending_objects:
                        pending_object = self._handler.pending_objects.pop(0)
                        count += 1
                        yield (pending_object, scope_url)

            except urllib.error.URLError as e:
                _log.warning('Fetching %s: %s', scope_url, e)
            except SAXParseException as e:
                _log.warning('Parsing %s: %s', scope_url, e)
            finally:
                try:
                    parser.close()
                except SAXParseException:
                    # Received malformed XML, such as XML with missing
                    # closing tags.  This is likely caused by a
                    # prematurely-terminated connection.
                    _log.warning('Parsing %s: incomplete scope list',
                                 scope_url)
                parser.reset()
        # Log successful completion
        _log.info('End of scope list (%d)', count)

    def get_count(self):
        '''Return our current understanding of the number of objects in
        scope.'''
        with self._lock:
            return self._handler.count class SessionVariables(object):
    '''A set of session variables.'''

    def __init__(self):
        self._vars = dict()
        self._lock = threading.Lock()
        self._between_get_and_set = False

    def filter_get(self, keys):
        '''Return a dict giving the total values of the variables listed in
        keys.'''
        ret = dict()
        with self._lock:
            for key in keys:
                if key in self._vars:
                    ret[key] = self._vars[key].filter_get()
                else:
                    ret[key] = 0.0
        return ret

    def filter_update(self, values):
        '''Add new values produced by a filter into the specified variables.
        @values is a map of keys and the quantities to add to the
        corresponding values.'''
        with self._lock:
            for key, value in values.items():
                var = self._vars.setdefault(key, _SessionVariable())
                var.filter_update(value, self._between_get_and_set)

    def client_get(self):
        '''Atomically engage the between_get_and_set interlock and return
        a dict of all values.'''
        ret = dict()
        with self._lock:
            self._between_get_and_set = True
            for key, var in self._vars.items():
                ret[key] = var.client_get()
        return ret

    def client_set(self, values):
        '''Atomically release the between_get_and_set interlock and update
        the session variables from the other servers.'''
        with self._lock:
            self._between_get_and_set = False
            for key, value in values.items():
                var = self._vars.setdefault(key, _SessionVariable())
                var.client_set(value) class SearchStatistics(_Statistics):
    label = 'Search statistics'
    attrs = (('objs_processed', 'Objects considered'),
             ('objs_dropped', 'Objects dropped'),
             ('objs_passed', 'Objects passed'),
             ('objs_true_positive', 'Objects passed with label'),
             ('objs_false_negative', 'Objects dropped with label'),
             ('objs_unloadable', 'Objects failing to load'),
             ('execution_us', 'Total object examination time (us)'))

    def __init__(self):
        super(SearchStatistics, self).__init__()

    def xdr(self, objs_total, filter_stats):
        '''Return an XDR statistics structure for these statistics.'''
        with self.lock:
            try:
                avg_obj_us = old_div(self.execution_us, self.objs_processed)
            except ZeroDivisionError:
                avg_obj_us = 0

            stats = []
            stats.append(XDR_stat('objs_total', objs_total))
            stats.append(XDR_stat('avg_obj_time_us', avg_obj_us))
            for name, _desc in self.attrs:
                if name != 'execution_us':
                    stats.append(XDR_stat(name, getattr(self, name)))

            return XDR_search_stats(
                stats=stats,
                filter_stats=[s.xdr() for s in filter_stats],
            ) class ResourceContext(object):
    """A resource context is a "namespace" for special resources.
    Resource can be: docker containers, etc.
    All resource must be created within a ResourceContext.
    When a ResourceContext is cleaned up, all resources therein must be
    cleaned up properly. All creations/modifications of resource therein
    are synchronized with a lock.
    ResourceContext can be shared by multiple processes if initialized
    with shared objects from the multiprocessing module"""

    def __init__(self, name, config, lock=threading.Lock(), catalog=dict()):
        super(ResourceContext, self).__init__()
        assert isinstance(config, DiamondConfig)
        self._name = name
        self._config = config
        self._lock = lock
        self._catalog = catalog  # resource signature -> handle (dict)
        _log.info('Created resource context %s', self._name)

    def ensure_resource(self, rtype, *args, **kargs):
        """Ensure a resource exists in the context. If already exists, return the handle.
        If not, create it and return the URI."""
        sig = _ResourceFactory.get_signature(rtype, *args, **kargs)
        with self._lock:
            if sig not in self._catalog: 
                for h in _ResourceFactory.create(rtype, self._config, *args, **kargs):
                    self._catalog[sig] = h
                    _log.debug("Update resource %s => %s", sig, str(h))
            rv = self._catalog[sig]

        return rv

    def cleanup(self):
        _log.info('Cleaning up resource context %s', self._name)
        with self._lock:
            for (_, h) in list(self._catalog.items()):
                _log.debug('Cleaning up %s', str(h))
                _ResourceFactory.cleanup(h)
            self._catalog.clear()
from builtins import str
from builtins import object
from functools import wraps
from opendiamond import protocol
from opendiamond.blobcache import ExecutableBlobCache
from opendiamond.protocol import (
    DiamondRPCFailure, DiamondRPCFCacheMiss, DiamondRPCCookieExpired,
    DiamondRPCSchemeNotSupported)
from opendiamond.rpc import RPCHandlers, RPCError, RPCProcedureUnavailable
from opendiamond.scope import ScopeCookie, ScopeError, ScopeCookieExpired
from opendiamond.server.filter import (
    FilterStack, Filter, FilterDependencyError, FilterUnsupportedSource)
from opendiamond.server.object_ import EmptyObject, Object, ObjectLoader
from opendiamond.server.scopelist import ScopeListLoader
from opendiamond.server.sessionvars import SessionVariables
from opendiamond.server.statistics import SearchStatistics
from opendiamond.server.resource import ResourceContext
import logging
import multiprocessing as mp
import os
import signal
import threading
        while self._workers:
            p = self._workers.pop()
            _log.debug("Terminating worker process %d", p.pid)
            try:
                p.terminate()   # send SIGTERM, which will be caught as _Signalled in the child
                p.join(1)   # allow it to clean up
                os.kill(p.pid, signal.SIGKILL)  # shoot it in the head
                p.join()
            except OSError: # perhaps process already dead
                pass

        # Log search statistics
        if self._running:
            self._state.stats.log()
            for filter in self._filters:
                filter.stats.log()

    # This is not a static method: it's only called when initializing the
    # class, and the staticmethod() decorator does not create a callable.
    # Also avoid complaints about accesses to self._running
    # pylint: disable=no-self-argument,protected-access
    def running(should_be_running):
        '''Decorator that specifies that the handler can only be called
        before, or after, the search has started running.'''

        def decorator(func):
            @wraps(func)
            def wrapper(self, *args, **kwargs):
                if self._running != should_be_running:
                    raise RPCProcedureUnavailable()
                return func(self, *args, **kwargs)

            return wrapper

        return decorator

    # pylint: enable=no-self-argument,protected-access

    def _check_runnable(self):
        '''Validate state preparatory to starting a search or reexecution.'''
        if self._scope is None:
            raise DiamondRPCFailure('No search scope configured')
        if not self._filters:
            _log.warning('No filters configured')
        # Ensure we have all filter code and blob arguments
        try:
            for filter in self._filters:
                filter.resolve(self._state)
        except FilterDependencyError as e:
            raise DiamondRPCFCacheMiss(str(e))

    @RPCHandlers.handler(25, protocol.XDR_setup, protocol.XDR_blob_list)
    @running(False)
    def setup(self, params):
        '''Configure the search and return a list of SHA256 signatures not
        present in the blob cache.'''

        def log_header(desc):
            _log.info('  %s:', desc)

        def log_item(key, fmt, *args):
            _log.info('    %-14s ' + fmt, key + ':', *args)

        # Create filter stack
        filters = []
        missing = set()
        _log.info('Filters:')
        for f in params.filters:
            unsupported = False
            try:
                if not Filter.source_available(self._state, f.code):
                    missing.add(f.code)
                    code_state = 'not cached'
                else:
                    code_state = 'cached'
            except FilterUnsupportedSource:
                unsupported = True
                code_state = 'unsupported'
            try:
                if not Filter.source_available(self._state, f.blob):
                    missing.add(f.blob)
                    blob_state = 'not cached'
                else:
                    blob_state = 'cached'
            except FilterUnsupportedSource:
                unsupported = True
                blob_state = 'unsupported'
            log_header(f.name)
            log_item('Code', '%s, %s', f.code, code_state)
            log_item('Blob', '%s, %s', f.blob, blob_state)
            log_item('Arguments', '%s', ', '.join(f.arguments) or '<none>')
            log_item('Dependencies', '%s',
                     ', '.join(f.dependencies) or '<none>')
            log_item('Minimum score', '%f', f.min_score)
            log_item('Maximum score', '%f', f.max_score)
            if unsupported:
                raise DiamondRPCSchemeNotSupported()
            filters.append(Filter(f.name, f.code, f.blob, f.min_score,
                                  f.max_score, f.arguments, f.dependencies))
        filterstack = FilterStack(filters)

        # Parse scope cookies
        try:
            cookies = [ScopeCookie.parse(c) for c in params.cookies]
            _log.info('Scope cookies:')
            for cookie in cookies:
                log_header(cookie.serial)
                log_item('Servers', '%s', ', '.join(cookie.servers))
                log_item('Scopes', '%s', ', '.join(cookie.scopeurls))
                log_item('Expires', '%s', cookie.expires)
                if self._state.config.security_cookie_no_verify:
                    _log.warn('Bypassing cookie verification.')
                else:
                    cookie.verify(self._state.config.serverids,
                                self._state.config.certdata)
            scope = ScopeListLoader(self._state.config, self._server_id,
                                    cookies)
        except ScopeCookieExpired as e:
            _log.warning('%s', e)
            raise DiamondRPCCookieExpired()

        except ScopeError as e: