func (c *CaptureApiHandler) getCaptureCount(r ApiResource) ApiResource {
capture := r.(*Capture)
tr := traversal.NewGremlinTraversalParser(strings.NewReader(capture.GremlinQuery), c.Graph)
ts, err := tr.Parse()
if err != nil {
logging.GetLogger().Errorf("Gremlin expression error: %s", err.Error())
return r
}
res, err := ts.Exec()
if err != nil {
logging.GetLogger().Errorf("Gremlin execution error: %s", err.Error())
return r
}
for _, value := range res.Values() {
switch value.(type) {
case *graph.Node:
n := value.(*graph.Node)
if common.IsCaptureAllowed(n.Metadata()["Type"].(string)) {
capture.Count = capture.Count + 1
}
case []*graph.Node:
capture.Count = capture.Count + len(value.([]*graph.Node))
default:
capture.Count = 0
}
}
return capture
}
func captureAllowedNodes(nodes []*graph.Node) []*graph.Node {
var allowed []*graph.Node
for _, n := range nodes {
if t, ok := n.Metadata()["Type"]; ok && common.IsCaptureAllowed(t.(string)) {
allowed = append(allowed, n)
}
}
return allowed
}
func (s *FlowGremlinTraversalStep) Exec(last traversal.GraphTraversalStep) (traversal.GraphTraversalStep, error) {
var graphTraversal *traversal.GraphTraversal
var err error
var paramsFilter *flow.Filter
if len(s.context.Params) > 0 {
if paramsFilter, err = paramsToFilter(s.context.Params...); err != nil {
return nil, err
}
}
var interval *flow.Range
if s.context.StepContext.Range != nil {
interval = &flow.Range{From: 0, To: s.context.StepContext.Range[1]}
}
flowset := flow.NewFlowSet()
switch tv := last.(type) {
case *traversal.GraphTraversal:
graphTraversal = tv
if context := graphTraversal.Graph.GetContext(); context.Time != nil && s.Storage != nil {
var flows []*flow.Flow
if flows, err = storage.LookupFlows(s.Storage, context, paramsFilter, interval); err == nil {
flowset.Flows = append(flowset.Flows, flows...)
}
} else {
flowSearchQuery := &flow.FlowSearchQuery{
Filter: paramsFilter,
Range: interval,
Sort: s.context.StepContext.Sort,
Dedup: s.context.StepContext.Dedup,
}
flowset, err = s.TableClient.LookupFlows(flowSearchQuery)
}
if r := s.context.StepContext.Range; r != nil {
flowset.Slice(int(r[0]), int(r[1]))
}
case *traversal.GraphTraversalV:
graphTraversal = tv.GraphTraversal
hnmap := make(flow.HostNodeIDMap)
graphTraversal.Graph.RLock()
for _, v := range tv.Values() {
node := v.(*graph.Node)
if t, ok := node.Metadata()["Type"]; !ok || !common.IsCaptureAllowed(t.(string)) {
continue
}
hnmap[node.Host()] = append(hnmap[node.Host()], string(node.ID))
}
graphTraversal.Graph.RUnlock()
if context := graphTraversal.Graph.GetContext(); context.Time != nil && s.Storage != nil {
var flows []*flow.Flow
if flows, err = storage.LookupFlowsByNodes(s.Storage, context, hnmap, paramsFilter, interval); err == nil {
flowset.Flows = append(flowset.Flows, flows...)
}
} else {
flowSearchQuery := &flow.FlowSearchQuery{
Filter: paramsFilter,
Range: interval,
Sort: s.context.StepContext.Sort,
Dedup: s.context.StepContext.Dedup,
}
flowset, err = s.TableClient.LookupFlowsByNodes(hnmap, flowSearchQuery)
}
if r := s.context.StepContext.Range; r != nil {
flowset.Slice(int(r[0]), int(r[1]))
}
default:
return nil, traversal.ExecutionError
}
if err != nil {
logging.GetLogger().Errorf("Error while looking for flows: %s", err.Error())
return nil, err
}
return &FlowTraversalStep{GraphTraversal: graphTraversal, flowset: flowset}, nil
}