def __init__(self, query=None, explain=None, columns=None, sort=None, width=None, filter=None, stream_results=False, complete_results=True, **kwargs): super(EntityFind, self).__init__(**kwargs) if query: self.query = entity_query.Query(query) else: self.query = entity_query.Query(self.search) if columns is not None: self.columns = columns if sort is not None: self.sort = sort if width is not None: self.width = width if filter is not None: self.display_filter = entity_query.Query(filter) if stream_results is not None: if sort: raise ValueError( "Can't stream results and sort at the same time.") self.stream_results = stream_results if complete_results is not None: self.complete_results = complete_results self.explain = explain
def find_by_component(self, component, complete=True): """Finds all entities that have the component. Arguments: complete: If True, will attempt to collect the component. """ query = entity_query.Query(expression.ComponentLiteral(component)) if complete: self.collect_for(query) return list(self.lookup_tables["components"].lookup(component))
def stream(self, query, handler, query_params=None): query = entity_query.Query(query, params=query_params) seen = set() def _deduplicator(entity): if entity in seen: return seen.add(entity) handler(entity) self.collect_for(query, result_stream_handler=_deduplicator) for entity in self.find(query, complete=False): _deduplicator(entity)
def get_referencing_entities(self, key, complete=True): """Finds entities that reference this entity by its identity. If other entities have an attribute that stores the identity of this entity then this function will find those entities. For example, calling this on a process, one can find all the handles owned by the process by calling process.get_referencing_entities("Handle/process"). Arguments: key: The property path to the attribute on the other entities. As usual, form is Component/attribute. """ # Automatically ask the entity manager to add indexing for # identity-based attributes. This is a good heuristic for optimal # performance. self.manager.add_attribute_lookup(key) return self.manager.find(query.Query("%s is {}" % key, params=[self.identity]), complete=complete)
def find(self, query, complete=True, validate=True, query_params=None): """Runs the query and yields entities that match. Arguments: query: Either an instance of the query AST, a query string, or a dictionary of queries. If a dict is given, a new dict will be returned with the same keys and values replaced with results. complete: If True, will trigger collectors as necessary, to ensure completness of results. validate: Will cause the query to be validated first (mostly for type errors.) """ if isinstance(query, dict): results = {} for query_name, expr in query.iteritems(): results[query_name] = self.find(expr, complete=complete, validate=validate, query_params=query_params) return results if not isinstance(query, entity_query.Query): query = entity_query.Query(query, params=query_params) if validate: query.execute("QueryValidator") if complete: self.collect_for(query) # Try to satisfy the query using available lookup tables. search = entity_lookup.EntityQuerySearch(query) return search.search(self.entities, self.lookup_tables)
def assertQueryMatches(self, query, bindings): m = entity_query.Query(query).execute("QueryMatcher", "match", bindings=bindings) self.assertIsNotNone(m)
def analyze(self, wanted): """Finds collectors and indexing suggestions for the query. Returns a dict of: - collectors: list of collectors to run - lookups: list of attributes to consider indexing for - dependencies: list of SimpleDependency instances to include - exclusions: list of SimpleDependency instances to exclude """ if not isinstance(wanted, entity_query.Query): wanted = entity_query.Query(wanted) # We cache by the source and not the query because we want to reanalyze # queries that are logically equivalent, but expressed differently, in # order to have the right cursor positions stored for highlighting in # GUI. cache_key = wanted.source analysis = self._cached_query_analyses.get(cache_key, None) if analysis: # We want to make a copy exactly one level deep. analysis_copy = {} for key, value in analysis.iteritems(): analysis_copy[key] = copy.copy(value) return analysis_copy analyzer = wanted.execute("QueryAnalyzer") include = analyzer.include exclude = analyzer.exclude suggested_indices = analyzer.latest_indices # A collector is a match if any of its promises match any of the # dependencies of the query. matched_collectors = [] for collector in self.collectors.itervalues(): for promise, dependency in itertools.product( collector.promises, include): if dependency.match(promise): matched_collectors.append(collector) break # A collector is yielded unless each one of its promises matches # an exclusion from dependencies. collectors = set() for collector in matched_collectors: for promise, exclusion in itertools.product( collector.promises, exclude): if not exclusion.match(promise): collectors.add(collector) break else: # No exclusions. collectors.add(collector) # A component is guaranteed if any dependency lists it. It is likely # if collectors we depend on output it (though not guaranteed). guaranteed_components = set(analyzer.expected_components) possible_components = set() for dependency in include: component = dependency.component if component in guaranteed_components: continue possible_components.add(dependency.component) for collector in collectors: for promise in collector.promises: component = promise.component if component in guaranteed_components: continue possible_components.add(component) analysis = dict(collectors=list(collectors), lookups=suggested_indices, dependencies=include, exclusions=exclude, guaranteed_components=guaranteed_components, possible_components=possible_components) self._cached_query_analyses[cache_key] = analysis return analysis
def _ensure_compile_queries(self): if self.collect_queries or self.collect_args is None: return for arg, source in self.collect_args.iteritems(): self.collect_queries[arg] = entity_query.Query(source)
def __init__(self, query=None, **kwargs): super(EntityAnalyze, self).__init__(**kwargs) self.query = entity_query.Query(query) self.analysis = self.session.entities.analyze(self.query)
def assertQueryRaises(self, query): q = entity_query.Query(query) self.assertRaises(validator.ValidationError, q.execute, "QueryValidator")
def assertQueryPasses(self, query): q = entity_query.Query(query) try: q.execute("QueryValidator") except validator.ValidationError as e: self.fail("%s raised exception on validation:\n%s" % (q, e))
def _subquery(self, expr): return entity_query.Query(expression=expr, source=self.query.source)
def analyze(self, query): return entity_query.Query(query).execute("QueryAnalyzer")