1// Copyright 2012-present Oliver Eilhard. All rights reserved.
2// Use of this source code is governed by a MIT-license.
3// See http://olivere.mit-license.org/license.txt for details.
4
5package elastic
6
7// IPRangeAggregation is a range aggregation that is dedicated for
8// IP addresses.
9//
10// See: https://www.elastic.co/guide/en/elasticsearch/reference/6.7/search-aggregations-bucket-iprange-aggregation.html
11type IPRangeAggregation struct {
12	field           string
13	subAggregations map[string]Aggregation
14	meta            map[string]interface{}
15	keyed           *bool
16	entries         []IPRangeAggregationEntry
17}
18
19type IPRangeAggregationEntry struct {
20	Key  string
21	Mask string
22	From string
23	To   string
24}
25
26func NewIPRangeAggregation() *IPRangeAggregation {
27	return &IPRangeAggregation{
28		subAggregations: make(map[string]Aggregation),
29		entries:         make([]IPRangeAggregationEntry, 0),
30	}
31}
32
33func (a *IPRangeAggregation) Field(field string) *IPRangeAggregation {
34	a.field = field
35	return a
36}
37
38func (a *IPRangeAggregation) SubAggregation(name string, subAggregation Aggregation) *IPRangeAggregation {
39	a.subAggregations[name] = subAggregation
40	return a
41}
42
43// Meta sets the meta data to be included in the aggregation response.
44func (a *IPRangeAggregation) Meta(metaData map[string]interface{}) *IPRangeAggregation {
45	a.meta = metaData
46	return a
47}
48
49func (a *IPRangeAggregation) Keyed(keyed bool) *IPRangeAggregation {
50	a.keyed = &keyed
51	return a
52}
53
54func (a *IPRangeAggregation) AddMaskRange(mask string) *IPRangeAggregation {
55	a.entries = append(a.entries, IPRangeAggregationEntry{Mask: mask})
56	return a
57}
58
59func (a *IPRangeAggregation) AddMaskRangeWithKey(key, mask string) *IPRangeAggregation {
60	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, Mask: mask})
61	return a
62}
63
64func (a *IPRangeAggregation) AddRange(from, to string) *IPRangeAggregation {
65	a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: to})
66	return a
67}
68
69func (a *IPRangeAggregation) AddRangeWithKey(key, from, to string) *IPRangeAggregation {
70	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: to})
71	return a
72}
73
74func (a *IPRangeAggregation) AddUnboundedTo(from string) *IPRangeAggregation {
75	a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: ""})
76	return a
77}
78
79func (a *IPRangeAggregation) AddUnboundedToWithKey(key, from string) *IPRangeAggregation {
80	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: ""})
81	return a
82}
83
84func (a *IPRangeAggregation) AddUnboundedFrom(to string) *IPRangeAggregation {
85	a.entries = append(a.entries, IPRangeAggregationEntry{From: "", To: to})
86	return a
87}
88
89func (a *IPRangeAggregation) AddUnboundedFromWithKey(key, to string) *IPRangeAggregation {
90	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: "", To: to})
91	return a
92}
93
94func (a *IPRangeAggregation) Lt(to string) *IPRangeAggregation {
95	a.entries = append(a.entries, IPRangeAggregationEntry{From: "", To: to})
96	return a
97}
98
99func (a *IPRangeAggregation) LtWithKey(key, to string) *IPRangeAggregation {
100	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: "", To: to})
101	return a
102}
103
104func (a *IPRangeAggregation) Between(from, to string) *IPRangeAggregation {
105	a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: to})
106	return a
107}
108
109func (a *IPRangeAggregation) BetweenWithKey(key, from, to string) *IPRangeAggregation {
110	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: to})
111	return a
112}
113
114func (a *IPRangeAggregation) Gt(from string) *IPRangeAggregation {
115	a.entries = append(a.entries, IPRangeAggregationEntry{From: from, To: ""})
116	return a
117}
118
119func (a *IPRangeAggregation) GtWithKey(key, from string) *IPRangeAggregation {
120	a.entries = append(a.entries, IPRangeAggregationEntry{Key: key, From: from, To: ""})
121	return a
122}
123
124func (a *IPRangeAggregation) Source() (interface{}, error) {
125	// Example:
126	// {
127	//     "aggs" : {
128	//         "range" : {
129	//             "ip_range": {
130	//                 "field": "ip",
131	//                 "ranges": [
132	//                     { "to": "10.0.0.5" },
133	//                     { "from": "10.0.0.5" }
134	//                 ]
135	//             }
136	//         }
137	//         }
138	//     }
139	// }
140	//
141	// This method returns only the { "ip_range" : { ... } } part.
142
143	source := make(map[string]interface{})
144	opts := make(map[string]interface{})
145	source["ip_range"] = opts
146
147	// ValuesSourceAggregationBuilder
148	if a.field != "" {
149		opts["field"] = a.field
150	}
151
152	if a.keyed != nil {
153		opts["keyed"] = *a.keyed
154	}
155
156	var ranges []interface{}
157	for _, ent := range a.entries {
158		r := make(map[string]interface{})
159		if ent.Key != "" {
160			r["key"] = ent.Key
161		}
162		if ent.Mask != "" {
163			r["mask"] = ent.Mask
164		} else {
165			if ent.From != "" {
166				r["from"] = ent.From
167			}
168			if ent.To != "" {
169				r["to"] = ent.To
170			}
171		}
172		ranges = append(ranges, r)
173	}
174	opts["ranges"] = ranges
175
176	// AggregationBuilder (SubAggregations)
177	if len(a.subAggregations) > 0 {
178		aggsMap := make(map[string]interface{})
179		source["aggregations"] = aggsMap
180		for name, aggregate := range a.subAggregations {
181			src, err := aggregate.Source()
182			if err != nil {
183				return nil, err
184			}
185			aggsMap[name] = src
186		}
187	}
188
189	// Add Meta data if available
190	if len(a.meta) > 0 {
191		source["meta"] = a.meta
192	}
193
194	return source, nil
195}
196