Rework complete project

pkg/downloader/downloader.go

@@ -0,0 +1,198 @@
+package downloader
+
+import (
+	"crypto/md5"
+	"crypto/sha1"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+	"hash"
+	"io"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/dustin/go-humanize"
+	req "github.com/levigross/grequests"
+)
+
+var (
+	hashes = map[string]hash.Hash{
+		"md5":    md5.New(),  //nolint: gosec
+		"sha1":   sha1.New(), //nolint: gosec
+		"sha256": sha256.New(),
+	}
+)
+
+type Context struct {
+	Filename   string
+	Filesize   int64
+	Reader     io.Reader
+	Chan       chan *Response
+	Closer     []func() error
+	hashWriter hash.Hash
+}
+
+func NewContext(filename string, filesize int64, hw hash.Hash, closer ...func() error) *Context {
+	return &Context{
+		Filename:   filename,
+		Filesize:   filesize,
+		Chan:       make(chan *Response),
+		Closer:     closer,
+		hashWriter: hw,
+	}
+}
+
+func (ctx *Context) Checksum() string {
+	return fmt.Sprintf("%x", ctx.hashWriter.Sum(nil))
+}
+
+func (ctx *Context) SetReader(rdr io.Reader) {
+	ctx.Reader = io.TeeReader(
+		rdr,
+		io.MultiWriter(
+			ctx.hashWriter,
+			&ProgressWriter{
+				Filename: ctx.Filename,
+				Filesize: ctx.Filesize,
+				Channel:  ctx.Chan,
+				Start:    time.Now(),
+			},
+		),
+	)
+}
+
+func (ctx *Context) Close() {
+	for i := len(ctx.Closer) - 1; i >= 0; i-- {
+		ctx.Closer[i]()
+	}
+}
+
+type ProgressWriter struct {
+	Channel  chan *Response
+	Start    time.Time
+	Filename string
+	Filesize int64
+	Bytes    int64
+}
+
+func (pw *ProgressWriter) Write(data []byte) (int, error) {
+	pw.Bytes = pw.Bytes + int64(len(data))
+	var progress float64
+	if pw.Filesize > 0 {
+		progress = float64(pw.Bytes) / float64(pw.Filesize) * 100
+	}
+	pw.Channel <- &Response{
+		Progress:        progress,
+		BytesPerSecond:  float64(pw.Bytes) / time.Since(pw.Start).Seconds(),
+		BytesTransfered: pw.Bytes,
+		BytesTotal:      pw.Filesize,
+	}
+	return len(data), nil
+}
+
+type File struct {
+	URL      string
+	Filename string
+	Size     int64
+}
+
+type Response struct {
+	Error           error
+	Progress        float64
+	BytesPerSecond  float64
+	BytesTransfered int64
+	BytesTotal      int64
+}
+
+func handleError(err error, adds ...string) {
+	msg := err.Error()
+	if len(adds) != 0 {
+		msg = fmt.Sprintf("%s %s", msg, strings.Join(adds, " "))
+	}
+	fmt.Println(msg)
+	os.Exit(0)
+}
+
+func printProgress(ctx *Context) {
+	for {
+		resp, ok := <-ctx.Chan
+		if !ok {
+			fmt.Println("\nSaved file to", ctx.Filename)
+			for key, value := range hashes {
+				fmt.Printf("%s: %x\n", key, value.Sum(nil))
+			}
+			break
+		}
+		if resp.Error != nil {
+			handleError(resp.Error)
+		}
+		if resp.BytesTotal > 0 {
+			fmt.Printf("\rDownloading %7s/%7s %.02f%%", humanize.Bytes(uint64(resp.BytesTransfered)), humanize.Bytes(uint64(resp.BytesTotal)), resp.Progress)
+		} else {
+			fmt.Printf("\rDownloading %7s", humanize.Bytes(uint64(resp.BytesTransfered)))
+		}
+	}
+}
+
+func DownloadFile(path, url string) {
+	file := File{
+		URL:      url,
+		Filename: path,
+	}
+	response, err := req.Head(file.URL, nil)
+	if err != nil {
+		if strings.Contains(err.Error(), "no such host") {
+			handleError(errors.New("invalid url"))
+		} else if strings.Contains(err.Error(), "certificate signed by unknown authority") {
+			handleError(errors.New("certificate signed by unknown authority"))
+		}
+		handleError(err)
+	}
+	defer response.Close()
+	if response.StatusCode == 404 {
+		handleError(errors.New("invalid url"))
+	}
+	if response.StatusCode == 401 {
+		handleError(errors.New("restriced access, credentials required"))
+	}
+	response.Close()
+
+	resp, err := req.Get(file.URL, nil)
+	if err != nil {
+		handleError(err)
+	}
+
+	var filesize int64
+	if resp.RawResponse.ContentLength != -1 {
+		filesize = resp.RawResponse.ContentLength
+	}
+	ctx := NewContext(file.Filename, filesize, md5.New())
+	ctx.SetReader(resp.RawResponse.Body)
+
+	go func(ctx *Context) {
+		defer ctx.Close()
+
+		dstFile, err := os.Create(ctx.Filename)
+		if err != nil {
+			ctx.Chan <- &Response{Error: err}
+			close(ctx.Chan)
+			return
+		}
+		defer dstFile.Close()
+
+		writer := io.MultiWriter(dstFile, hashes["md5"], hashes["sha1"], hashes["sha256"])
+		_, err = io.Copy(writer, ctx.Reader)
+		if err != nil {
+			ctx.Chan <- &Response{Error: err}
+		}
+		close(ctx.Chan)
+	}(ctx)
+
+	printProgress(ctx)
+}
+
+func FileExists(path string) bool {
+	_, err := os.Stat(path)
+	return !errors.Is(err, os.ErrNotExist)
+}

pkg/geoloc/csv_import.go

@@ -0,0 +1,79 @@
+package geoloc
+
+import (
+	"encoding/csv"
+	"os"
+	"strconv"
+)
+
+func parseCSV(filePath string) ([]IPInfo, error) {
+	file, err := os.Open(filePath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	reader := csv.NewReader(file)
+
+	var data []IPInfo
+	for {
+		record, err := reader.Read()
+		if err != nil {
+			break
+		}
+
+		ipnumfrom, err := strconv.Atoi(record[0])
+		if err != nil {
+			continue
+		}
+		ipnumto, err := strconv.Atoi(record[1])
+		if err != nil {
+			continue
+		}
+		if record[2] == "-" {
+			record[2] = ""
+		}
+		if record[3] == "-" {
+			record[3] = ""
+		}
+		if record[4] == "-" {
+			record[4] = ""
+		}
+		if record[5] == "-" {
+			record[5] = ""
+		}
+		latitude, err := strconv.ParseFloat(record[6], 32)
+		if err != nil {
+			latitude = 0
+		}
+		longitude, err := strconv.ParseFloat(record[7], 32)
+		if err != nil {
+			longitude = 0
+		}
+
+		ipinfo := IPInfo{
+			IPNumFrom: uint(ipnumfrom),
+			IPNumTo:   uint(ipnumto),
+			Code:      record[2],
+			Country:   record[3],
+			State:     record[4],
+			City:      record[5],
+			Latitude:  float32(latitude),
+			Longitude: float32(longitude),
+		}
+
+		data = append(data, ipinfo)
+	}
+
+	return data, nil
+}
+
+func ImportCSV(filePath string) error {
+	ipinfos, err := parseCSV(filePath)
+	if err != nil {
+		return err
+	}
+
+	err = CreateDatabase(ipinfos)
+	return err
+}

pkg/geoloc/database.go

@@ -0,0 +1,75 @@
+package geoloc
+
+import (
+	"fmt"
+
+	"github.com/hashicorp/go-memdb"
+)
+
+var database *memdb.MemDB
+
+func CreateDatabase(ipinfos []IPInfo) error {
+	schema := &memdb.DBSchema{
+		Tables: map[string]*memdb.TableSchema{
+			"ipinfo": {
+				Name: "ipinfo",
+				Indexes: map[string]*memdb.IndexSchema{
+					"id": {
+						Name:    "id",
+						Unique:  true,
+						Indexer: &memdb.UintFieldIndex{Field: "IPNumFrom"},
+					},
+				},
+			},
+		},
+	}
+
+	db, err := memdb.NewMemDB(schema)
+	if err != nil {
+		return err
+	}
+
+	txn := db.Txn(true)
+	defer txn.Abort()
+
+	for _, ipinfo := range ipinfos {
+		if err := txn.Insert("ipinfo", ipinfo); err != nil {
+			return err
+		}
+	}
+
+	txn.Commit()
+
+	database = db
+	return nil
+}
+
+func SearchIPNet(ipnetnum uint) (*IPInfo, error) {
+	txn := database.Txn(false)
+	defer txn.Abort()
+
+	raw, err := txn.First("ipinfo", "id", ipnetnum)
+	if err != nil {
+		return nil, err
+	}
+	if raw == nil {
+		return nil, nil
+	}
+
+	var ipinfo IPInfo
+	for {
+		raw, err := txn.First("ipinfo", "id", ipnetnum)
+		if err != nil {
+			return nil, err
+		}
+		if raw != nil {
+			ipinfo = raw.(IPInfo)
+			break
+		}
+		if ipnetnum == 0 {
+			return nil, fmt.Errorf("SearchIPNet: IP net not found")
+		}
+		ipnetnum = ipnetnum - 256
+	}
+	return &ipinfo, nil
+}

pkg/geoloc/main.go

@@ -1,10 +0,0 @@
-package geoloc
-
-import "strings"
-
-func GetCountry(input string) *Country {
-	if c, exists := countries[strings.ToUpper(input)]; exists {
-		return &c
-	}
-	return nil
-}

pkg/geoloc/static.go

@@ -0,0 +1,30 @@
+package geoloc
+
+import (
+	"fmt"
+	"math"
+	"net"
+	"strings"
+
+	"github.com/praserx/ipconv"
+)
+
+func GetCountry(input string) *IPInfo {
+	if c, exists := countries[strings.ToUpper(input)]; exists {
+		return &c
+	}
+	return nil
+}
+
+func GetIPInfo(ipaddress string) (uint, error) {
+	ip := net.ParseIP(ipaddress)
+	if ip == nil {
+		return 0, fmt.Errorf("no valid IP address")
+	}
+	ipnum, err := ipconv.IPv4ToInt(ip)
+	if err != nil {
+		return 0, err
+	}
+	ipnumfrom := uint(math.Floor(float64(ipnum)/float64(256))) * 256
+	return ipnumfrom, nil
+}

pkg/geoloc/static_test.go

@@ -12,7 +12,7 @@ func TestGetCountry(t *testing.T) {
 	tests := []struct {
 		name string
 		args args
-		want *Country
+		want *IPInfo
 	}{
 		{"GBR Test", args{input: "826"}, nil},
 		{"USA upper Test", args{input: "USA"}, nil},
@@ -31,3 +31,23 @@ func TestGetCountry(t *testing.T) {
 		})
 	}
 }
+
+func TestGetIPInfo(t *testing.T) {
+	type args struct {
+		input string
+	}
+	tests := []struct {
+		name string
+		args args
+		want uint
+	}{
+		{"IP address Test", args{input: "1.1.1.1"}, 16843008},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got, err := GetIPInfo(tt.args.input); err != nil || !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GetIPInfo() - %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

pkg/geoloc/types.go

@@ -0,0 +1,12 @@
+package geoloc
+
+type IPInfo struct {
+	IPNumFrom uint    `json:"ip_num_min"`
+	IPNumTo   uint    `json:"ip_num_max"`
+	Code      string  `json:"code"`
+	Country   string  `json:"country"`
+	State     string  `json:"state"`
+	City      string  `json:"city"`
+	Latitude  float32 `json:"latitude"`
+	Longitude float32 `json:"longitude"`
+}