proxy_util/main.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"time"

	log "github.com/fangdingjun/go-log/v5"
	"gopkg.in/yaml.v2"
)

var iplist = []string{}

// Conf config file
type Conf struct {
	Port     int    `yaml:"port"`
	ListFile string `yaml:"ip_list"`
}

func main() {
	var cfg Conf
	var cfgfile string
	var logfile string
	var loglevel string

	flag.StringVar(&cfgfile, "c", "config.yaml", "config file")
	flag.StringVar(&logfile, "log_file", "", "log file")
	flag.StringVar(&loglevel, "log_level", "INFO", "log level")
	flag.Parse()

	lv, err := log.ParseLevel(loglevel)
	if err != nil {
		fmt.Printf("wrong level name: %s\nvalid level name: DEBUG, INFO, WARN, ERROR, FATAL\n", err)
		os.Exit(1)
	}
	log.Default.Level = lv

	if logfile != "" {
		log.Default.Out = &log.FixedSizeFileWriter{
			Name:     logfile,
			MaxSize:  10 * 1024 * 1024,
			MaxCount: 4,
		}
	}

	rand.Seed(time.Now().Unix())

	data, err := ioutil.ReadFile(cfgfile)
	if err != nil {
		log.Fatal(err)
	}
	if err = yaml.Unmarshal(data, &cfg); err != nil {
		log.Fatal(err)
	}
	loadIPList(cfg.ListFile)

	log.Printf("Listen at http://0.0.0.0:%d", cfg.Port)

	if err = http.ListenAndServe(fmt.Sprintf(":%d", cfg.Port), &handler{}); err != nil {
		log.Fatal(err)
	}
}

func loadIPList(f string) {
	fp, err := os.Open(f)
	if err != nil {
		log.Errorln(err)
		return
	}
	defer fp.Close()

	r := bufio.NewReader(fp)
	for {
		line, err := r.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				log.Errorln(err)
			}
			break
		}
		line = strings.Trim(line, " \r\n\t")
		if line == "" || line[0] == '#' {
			continue
		}
		iplist = append(iplist, line)
	}
}