Merge pull request #22 from m-messiah/issue-21

Generate ip2proxy Pro base from local file

Author: m-messiah

.gitignore

@@ -4,3 +4,4 @@ output
 *.exe
 .vscode
 output/
+.DS_Store

app.go

@@ -11,8 +11,13 @@ func main() {
 	outputDir := flag.String("output", "output", "output directory for files")
 	ipgeobase := flag.Bool("ipgeobase", false, "enable ipgeobase generation")
 	tor := flag.Bool("tor", false, "enable tor generation")
-	ip2proxyFlag := flag.Bool("ip2proxy", false, "enable ip2proxy generation")
-	ip2proxyToken := flag.String("ip2proxy-token", "", "Get token here https://lite.ip2location.com/file-download")
+	ip2proxyLiteFlag := flag.Bool("ip2proxy", false, "enable ip2proxy PX4-LITE generation")
+	ip2proxyFlag := flag.Bool("ip2proxy-pro", false, "enable ip2proxy PX4 generation")
+	ip2proxyLiteToken := flag.String("ip2proxy-token", "", "Get token here https://lite.ip2location.com/file-download")
+	ip2proxyToken := flag.String("ip2proxy-pro-token", "", "ip2proxy download token")
+	ip2proxyLiteFilename := flag.String("ip2proxy-lite-filename", "", "Filename of already downloaded ip2proxy-lite db")
+	ip2proxyFilename := flag.String("ip2proxy-pro-filename", "", "Filename of already downloaded ip2proxy db")
+	ip2proxyPrintType := flag.Bool("ip2proxy-print-type", false, "Print proxy type in map, instead of `1`")
 	maxmind := flag.Bool("maxmind", false, "enable maxmind generation")
 	maxmindIPVer := flag.Int("ipver", 4, "MaxMind ip version (4 or 6)")
 	maxmindLang := flag.String("lang", "ru", "MaxMind city name language")
@@ -29,12 +34,13 @@ func main() {
 		printMessage("ip2geo", "version "+VERSION, "OK")
 		return
 	}
-	if !(*ipgeobase || *tor || *maxmind || *ip2proxyFlag) {
+	if !(*ipgeobase || *tor || *maxmind || *ip2proxyLiteFlag || *ip2proxyFlag) {
 		// By default, generate all maps
 		*ipgeobase = true
 		*tor = true
 		*maxmind = true
-		*ip2proxyFlag = *ip2proxyToken != ""
+		*ip2proxyLiteFlag = *ip2proxyLiteToken != "" || *ip2proxyLiteFilename != ""
+		*ip2proxyFlag = *ip2proxyToken != "" || *ip2proxyFilename != ""
 	}
 	if *quiet {
 		logLevel = 1
@@ -82,12 +88,28 @@ func main() {
 		go Generate(&m)
 	}
 
+	if *ip2proxyLiteFlag {
+		goroutinesCount++
+		o := ip2proxy{
+			Name:       "ip2proxyLite",
+			Token:      *ip2proxyLiteToken,
+			Filename:   *ip2proxyLiteFilename,
+			ErrorsChan: errorChannel,
+			OutputDir:  *outputDir,
+			PrintType:  *ip2proxyPrintType,
+		}
+		go o.Get()
+	}
+
 	if *ip2proxyFlag {
 		goroutinesCount++
 		o := ip2proxy{
+			Name:       "ip2proxyPro",
 			Token:      *ip2proxyToken,
+			Filename:   *ip2proxyFilename,
 			ErrorsChan: errorChannel,
 			OutputDir:  *outputDir,
+			PrintType:  *ip2proxyPrintType,
 		}
 		go o.Get()
 	}

ip2proxy.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"archive/zip"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net"
@@ -23,43 +24,72 @@ type ip2proxyItem struct {
 }
 
 type ip2proxy struct {
-	items      []*ip2proxyItem
-	archive    []*zip.File
-	OutputDir  string
-	ErrorsChan chan Error
-	Token      string
+	items       []*ip2proxyItem
+	archive     []*zip.File
+	OutputDir   string
+	ErrorsChan  chan Error
+	Token       string
+	Filename    string
+	Name        string
+	csvFilename string
+	zipFilename string
+	PrintType   bool
 }
 
-func (o *ip2proxy) Get() {
-	answer, err := o.download()
+func (o *ip2proxy) checkErr(err error, message string) bool {
 	if err != nil {
-		o.ErrorsChan <- Error{err, "ip2proxy", "Download"}
+		o.ErrorsChan <- Error{err, o.Name, message}
+		return true
+	}
+	printMessage(o.Name, message, "OK")
+	return false
+}
+
+func (o *ip2proxy) Get() {
+	if o.Name == "ip2proxyPro" {
+		o.csvFilename = "IP2PROXY-IP-PROXYTYPE-COUNTRY-REGION-CITY-ISP.CSV"
+		o.zipFilename = "PX4-IP-PROXYTYPE-COUNTRY-REGION-CITY-ISP"
+	} else if o.Name == "ip2proxyLite" {
+		o.csvFilename = "IP2PROXY-LITE-PX4.CSV"
+		o.zipFilename = "PX4LITE"
+	} else {
+		o.ErrorsChan <- Error{errors.New("Unknown ip2proxy type requested"), o.Name, "bad init"}
 		return
 	}
-	printMessage("ip2proxy", "Download", "OK")
-	if err := o.unpack(answer); err != nil {
-		o.ErrorsChan <- Error{err, "ip2proxy", "Unpack"}
+	fileData, err := o.getZip()
+	if o.checkErr(err, "Get ZIP") {
 		return
 	}
-	printMessage("ip2proxy", "Unpack", "OK")
-	if err := o.Parse("IP2PROXY-LITE-PX4.CSV"); err != nil {
-		o.ErrorsChan <- Error{err, "ip2proxy", "Parse"}
+	err = o.unpack(fileData)
+	if o.checkErr(err, "Unpack") {
 		return
 	}
-	printMessage("ip2proxy", "Parse", "OK")
-	if err := o.Write(); err != nil {
-		o.ErrorsChan <- Error{err, "ip2proxy", "Write Nginx Map"}
+	err = o.Parse(o.csvFilename)
+	if o.checkErr(err, "Parse") {
+		return
+	}
+	err = o.Write()
+	if o.checkErr(err, "Write Nginx Map") {
 		return
 	}
-	printMessage("ip2proxy", "Write Nginx Map", "OK")
 	o.ErrorsChan <- Error{err: nil}
 }
 
+func (o *ip2proxy) getZip() ([]byte, error) {
+	if len(o.Token) > 0 {
+		return o.download()
+	} else if len(o.Filename) > 0 {
+		return ioutil.ReadFile(o.Filename)
+	} else {
+		return nil, errors.New("Token or Filename must be passed")
+	}
+}
+
 func (o *ip2proxy) download() ([]byte, error) {
 	client := &http.Client{}
 	req, err := http.NewRequest("GET", "https://www.ip2location.com/download", nil)
 	q := req.URL.Query()
-	q.Add("file", "PX4LITE")
+	q.Add("file", o.zipFilename)
 	q.Add("token", o.Token)
 	req.URL.RawQuery = q.Encode()
 	resp, err := client.Do(req)
@@ -87,10 +117,10 @@ func (o *ip2proxy) unpack(response []byte) error {
 
 func (o *ip2proxy) Parse(filename string) error {
 	var list []*ip2proxyItem
-	for record := range readCSVDatabase(o.archive, filename, "ip2proxy", ',', false) {
+	for record := range readCSVDatabase(o.archive, filename, o.Name, ',', false) {
 		item, err := o.lineToItem(record)
 		if err != nil {
-			printMessage("ip2proxy", fmt.Sprintf("Can't parse line from %s with %v", filename, err), "WARN")
+			printMessage(o.Name, fmt.Sprintf("Can't parse line from %s with %v", filename, err), "WARN")
 			continue
 		}
 		list = append(list, item)
@@ -130,13 +160,19 @@ func (o *ip2proxy) Write() error {
 }
 
 func (o *ip2proxy) writeNetworks() error {
-	file, err := os.Create(path.Join(o.OutputDir, "ip2proxy_net.txt"))
+	file, err := os.Create(path.Join(o.OutputDir, o.Name+"_net.txt"))
 	if err != nil {
 		return err
 	}
 	defer file.Close()
+	var mapValue string
 	for _, item := range o.items {
-		fmt.Fprintf(file, "%s-%s 1;\n", item.IPFrom, item.IPTo)
+		if o.PrintType {
+			mapValue = item.ProxyType
+		} else {
+			mapValue = "1"
+		}
+		fmt.Fprintf(file, "%s-%s \"%s\";\n", item.IPFrom, item.IPTo, mapValue)
 	}
 	return nil
 }

utils.go

@@ -62,7 +62,7 @@ func printMessage(module, message, status string) {
 		}
 		statusMesage = color.BlueString(status)
 	}
-	fmt.Printf("%-10s | %-60s [%s]\n", module, message, statusMesage)
+	fmt.Printf("%-15s | %-60s [%s]\n", module, message, statusMesage)
 }
 
 func getIPRange(ipver int, network string) string {