package main

import (
	"archive/tar"
	"compress/gzip"
	"github.com/oschwald/geoip2-golang"
	log "github.com/sirupsen/logrus"
	"io"
	"net"
	"os"
	"strings"
	"sync"
)

var geoLookupMutex = sync.Mutex{}
var city *geoip2.Reader
var asn *geoip2.Reader

func geoLookup(ip net.IP) (*geoip2.City, error) {
	geoLookupMutex.Lock()
	defer geoLookupMutex.Unlock()

	// check is GeoLite2-City.mmdb exists
	err := ensureDatabase("GeoLite2-City")
	if err != nil {
		return nil, err
	}

	if city == nil {
		var err error
		city, err = geoip2.Open("GeoLite2-City.mmdb")
		if err != nil {
			return nil, err
		}
	}

	// lookup
	record, err := city.City(ip)
	if err != nil {
		return nil, err
	}

	return record, nil
}

func asnLookup(ip net.IP) (*geoip2.ASN, error) {
	geoLookupMutex.Lock()
	defer geoLookupMutex.Unlock()

	// check is GeoLite2-City.mmdb exists
	err := ensureDatabase("GeoLite2-ASN")
	if err != nil {
		return nil, err
	}

	if asn == nil {
		var err error
		asn, err = geoip2.Open("GeoLite2-ASN.mmdb")
		if err != nil {
			return nil, err
		}
	}

	// lookup
	record, err := asn.ASN(ip)
	if err != nil {
		return nil, err
	}

	return record, nil
}

func ensureDatabase(databaseName string) error {
	if _, err := os.Stat(databaseName + ".mmdb"); os.IsNotExist(err) {
		// download
		url := "https://raw.githubusercontent.com/GitSquared/node-geolite2-redist/master/redist/" + databaseName + ".tar.gz"
		target := databaseName + ".tar.gz"
		err := downloadFile(url, target)
		if err != nil {
			return err
		}

		// extract
		_, err = extractDatabase(target, databaseName)
		if err != nil {
			return err
		}
	}
	return nil
}

func extractDatabase(source, databaseName string) (string, error) {
	tarFile, err := os.Open(source)
	if err != nil {
		return "", err
	}

	defer tarFile.Close()
	gzr, err := gzip.NewReader(tarFile)
	if err != nil {
		return "", err
	}

	defer gzr.Close()
	tr := tar.NewReader(gzr)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return "", err
		}
		log.Debugf("Extracting %s", header.Name)
		if strings.HasSuffix(header.Name, ".mmdb") {
			outFile, err := os.Create(databaseName + ".mmdb")
			if err != nil {
				return "", err
			}
			defer outFile.Close()
			_, err = io.Copy(outFile, tr)
			if err != nil {
				return "", err
			}
			break
		}
	}

	return "", nil
}