[feat] add prometheus metrics for country
[getmyip.git] / main.go
1 // Copyright (c) 2018 LEAP Encryption Access Project
2 //
3 // This program is free software: you can redistribute it and/or modify
4 // it under the terms of the GNU General Public License as published by
5 // the Free Software Foundation, either version 3 of the License, or
6 // (at your option) any later version.
7 //
8 // This program is distributed in the hope that it will be useful,
9 // but WITHOUT ANY WARRANTY; without even the implied warranty of
10 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 // GNU General Public License for more details.
12 //
13 // You should have received a copy of the GNU General Public License
14 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
15
16 package main
17
18 import (
19         "encoding/json"
20         "flag"
21         "fmt"
22         "log"
23         "math/rand"
24         "net"
25         "net/http"
26         "os"
27         "regexp"
28         "strconv"
29         "strings"
30         "time"
31
32         "github.com/StefanSchroeder/Golang-Ellipsoid/ellipsoid"
33         "github.com/hongshibao/go-kdtree"
34         "github.com/oschwald/geoip2-golang"
35         "github.com/prometheus/client_golang/prometheus"
36         "github.com/prometheus/client_golang/prometheus/promhttp"
37         "github.com/tidwall/cities"
38 )
39
40 func floatToString(num float64) string {
41         return strconv.FormatFloat(num, 'f', 6, 64)
42 }
43
44 func getRemoteIP(req *http.Request) string {
45         forward := req.Header.Get("X-Forwarded-For")
46         ipstr := ""
47         if forward != "" {
48                 ipstr = forward
49         } else {
50                 ip, _, err := net.SplitHostPort(req.RemoteAddr)
51                 if err != nil {
52                         log.Fatal(err)
53                 }
54                 netIP := net.ParseIP(ip)
55                 ipstr = netIP.String()
56         }
57         return ipstr
58 }
59
60 type geodb struct {
61         db          *geoip2.Reader
62         Forbidden   []string
63         Gateways    []gateway
64         GatewayTree *kdtree.KDTree
65         GatewayMap  map[[3]float64][]gateway
66         earth       *ellipsoid.Ellipsoid
67 }
68
69 func (g *geodb) getPointForLocation(lat float64, lon float64) *EuclideanPoint {
70         x, y, z := g.earth.ToECEF(lat, lon, 0)
71         p := NewEuclideanPoint(x, y, z)
72         return p
73 }
74
75 func randomizeGateways(gws []gateway) []gateway {
76         dest := make([]gateway, len(gws))
77         perm := rand.Perm(len(gws))
78         for i, v := range perm {
79                 dest[v] = gws[i]
80         }
81         return dest
82 }
83
84 func (g *geodb) sortGateways(lat float64, lon float64) []string {
85         ret := make([]string, 0)
86         t := g.getPointForLocation(lat, lon)
87         nn := g.GatewayTree.KNN(t, len(g.Gateways))
88         for i := 0; i < len(nn); i++ {
89                 p := [3]float64{nn[i].GetValue(0), nn[i].GetValue(1), nn[i].GetValue(2)}
90                 cityGateways := g.GatewayMap[p]
91                 if len(cityGateways) > 1 {
92                         cityGateways = randomizeGateways(cityGateways)
93                 }
94                 for _, gw := range cityGateways {
95                         if !stringInSlice(gw.Host, g.Forbidden) {
96                                 if !stringInSlice(gw.Host, ret) {
97                                         ret = append(ret, gw.Host)
98                                 }
99                         }
100                 }
101         }
102         return ret
103 }
104
105 func stringInSlice(a string, list []string) bool {
106         for _, b := range list {
107                 if b == a {
108                         return true
109                 }
110         }
111         return false
112 }
113
114 func (g *geodb) geolocateGateways(b *bonafide) {
115         g.GatewayMap = make(map[[3]float64][]gateway)
116         gatewayPoints := make([]kdtree.Point, 0)
117
118         for i := 0; i < len(b.eip.Gateways); i++ {
119                 gw := b.eip.Gateways[i]
120                 coord := geolocateCity(gw.Location)
121                 gw.Coordinates = coord
122                 b.eip.Gateways[i] = gw
123
124                 p := g.getPointForLocation(coord.Latitude, coord.Longitude)
125
126                 gatewayPoints = append(gatewayPoints, *p)
127                 var i [3]float64
128                 copy(i[:], p.Vec)
129                 g.GatewayMap[i] = append(g.GatewayMap[i], gw)
130         }
131         g.Gateways = b.eip.Gateways
132         g.GatewayTree = kdtree.NewKDTree(gatewayPoints)
133 }
134
135 func (g *geodb) getRecordForIP(ipstr string) *geoip2.City {
136         ip := net.ParseIP(ipstr)
137         record, err := g.db.City(ip)
138         if err != nil {
139                 log.Fatal(err)
140         }
141         return record
142 }
143
144 func geolocateCity(city string) coordinates {
145         // because some cities apparently are not good enough for the top 10k
146         missingCities := make(map[string]coordinates)
147         missingCities["hongkong"] = coordinates{22.319201099, 114.1696121}
148
149         re := regexp.MustCompile("-| ")
150         for _, c := range cities.Cities {
151                 canonical := strings.ToLower(city)
152                 canonical = re.ReplaceAllString(canonical, "")
153                 if strings.ToLower(c.City) == canonical {
154                         return coordinates{c.Latitude, c.Longitude}
155                 }
156                 v, ok := missingCities[canonical]
157                 if ok == true {
158                         return v
159                 }
160
161         }
162         return coordinates{0, 0}
163 }
164
165 type jsonHandler struct {
166         geoipdb *geodb
167 }
168
169 type GeolocationJSON struct {
170         Ip        string   `json:"ip"`
171         Cc        string   `json:"cc"`
172         City      string   `json:"city"`
173         Latitude  float64  `json:"lat"`
174         Longitude float64  `json:"lon"`
175         Gateways  []string `json:"gateways"`
176 }
177
178 func (jh *jsonHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
179         ipstr := getRemoteIP(req)
180         record := jh.geoipdb.getRecordForIP(ipstr)
181         sortedGateways := jh.geoipdb.sortGateways(record.Location.Latitude, record.Location.Longitude)
182
183         hitsPerCountry.With(prometheus.Labels{"country": record.Country.IsoCode}).Inc()
184
185         data := &GeolocationJSON{
186                 ipstr,
187                 record.Country.IsoCode,
188                 record.City.Names["en"],
189                 record.Location.Latitude,
190                 record.Location.Longitude,
191                 sortedGateways,
192         }
193
194         dataJSON, _ := json.Marshal(data)
195         fmt.Fprintf(w, string(dataJSON))
196 }
197
198 type txtHandler struct {
199         geoipdb *geodb
200 }
201
202 func (th *txtHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
203         ipstr := getRemoteIP(req)
204         record := th.geoipdb.getRecordForIP(ipstr)
205
206         fmt.Fprintf(w, "Your IP: %s\n", ipstr)
207         fmt.Fprintf(w, "Your Country: %s\n", record.Country.IsoCode)
208         fmt.Fprintf(w, "Your City: %s\n", record.City.Names["en"])
209         fmt.Fprintf(w, "Your Coordinates: %s, %s\n",
210                 floatToString(record.Location.Latitude),
211                 floatToString(record.Location.Longitude))
212 }
213
214 func main() {
215         rand.Seed(time.Now().UnixNano())
216         var port = flag.Int("port", 9001, "port where the service listens on")
217         var metricsPort = flag.Int("metricsPort", 9002, "port where the metrics server listens on")
218         var dbpath = flag.String("geodb", "/var/lib/GeoIP/GeoLite2-City.mmdb", "path to the GeoLite2-City database")
219         var notls = flag.Bool("notls", false, "disable TLS on the service")
220         var key = flag.String("server_key", "", "path to the key file for TLS")
221         var crt = flag.String("server_crt", "", "path to the cert file for TLS")
222         var forbidstr = flag.String("forbid", "", "comma-separated list of forbidden gateways")
223         flag.Parse()
224
225         forbidden := strings.Split(*forbidstr, ",")
226         fmt.Println("Forbidden gateways:", forbidden)
227
228         if *notls == false {
229                 if *key == "" || *crt == "" {
230                         log.Fatal("you must provide -server_key and -server_crt parameters")
231                 }
232                 if _, err := os.Stat(*crt); os.IsNotExist(err) {
233                         log.Fatal("path for crt file does not exist!")
234                 }
235                 if _, err := os.Stat(*key); os.IsNotExist(err) {
236                         log.Fatal("path for key file does not exist!")
237                 }
238         }
239
240         db, err := geoip2.Open(*dbpath)
241         if err != nil {
242                 log.Fatal(err)
243         }
244         defer db.Close()
245
246         earth := ellipsoid.Init("WGS84", ellipsoid.Degrees, ellipsoid.Meter, ellipsoid.LongitudeIsSymmetric, ellipsoid.BearingIsSymmetric)
247         geoipdb := geodb{db, forbidden, nil, nil, nil, &earth}
248
249         log.Println("Seeding gateway list...")
250         bonafide := newBonafide()
251         bonafide.getGateways()
252
253         geoipdb.geolocateGateways(bonafide)
254         bonafide.listGateways()
255
256         mux := http.NewServeMux()
257         jh := &jsonHandler{&geoipdb}
258         mux.Handle("/json", jh)
259
260         th := &txtHandler{&geoipdb}
261         mux.Handle("/", th)
262
263         mtr := http.NewServeMux()
264         mtr.Handle("/metrics", promhttp.Handler())
265
266         /* prometheus metrics */
267         go func() {
268                 pstr := ":" + strconv.Itoa(*metricsPort)
269                 log.Println("/metrics endpoint listening in port", *metricsPort)
270                 log.Fatal(http.ListenAndServe(pstr, mtr))
271         }()
272
273         /* geolocation api */
274         log.Println("Started Geolocation Service")
275         log.Printf("Listening on port %v...\n", *port)
276
277         pstr := ":" + strconv.Itoa(*port)
278         if *notls == true {
279                 err = http.ListenAndServe(pstr, mux)
280         } else {
281                 err = http.ListenAndServeTLS(pstr, *crt, *key, mux)
282         }
283
284         if err != nil {
285                 log.Fatal("error in listenAndServe[TLS]: ", err)
286         }
287 }