From 8a2181f822f8797770215e04685acf42a04be780 Mon Sep 17 00:00:00 2001
From: zu1k <i@lgf.im>
Date: Thu, 13 Aug 2020 08:18:52 +0800
Subject: [PATCH] auto fetch new config & add subs

---
 app/getter.go    | 15 +++++++++++++++
 app/task.go      | 28 ++++++++++++++++++++++++++++
 config/config.go |  2 +-
 main.go          | 18 +++++-------------
 4 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/app/getter.go b/app/getter.go
index 66ffdce..5c72d01 100644
--- a/app/getter.go
+++ b/app/getter.go
@@ -2,6 +2,7 @@ package app
 
 import (
 	"fmt"
+	"os"
 
 	"github.com/zu1k/proxypool/config"
 	"github.com/zu1k/proxypool/getter"
@@ -9,7 +10,21 @@ import (
 
 var Getters = make([]getter.Getter, 0)
 
+func InitConfigAndGetters(path string) {
+	c, err := config.Parse(path)
+	if err != nil {
+		fmt.Println("Error: ", err.Error())
+		os.Exit(1)
+	}
+	if c == nil {
+		fmt.Println("Error: no sources")
+		os.Exit(2)
+	}
+	InitGetters(c.Sources)
+}
+
 func InitGetters(sources []config.Source) {
+	Getters = make([]getter.Getter, 0)
 	for _, source := range sources {
 		g := getter.NewGetter(source.Type, source.Options)
 		if g != nil {
diff --git a/app/task.go b/app/task.go
index 7f65fbb..0d52e7c 100644
--- a/app/task.go
+++ b/app/task.go
@@ -1,17 +1,25 @@
 package app
 
 import (
+	"io/ioutil"
 	"log"
 	"math/rand"
+	"os"
 	"strconv"
 	"sync"
 
 	"github.com/zu1k/proxypool/app/cache"
 	"github.com/zu1k/proxypool/provider"
 	"github.com/zu1k/proxypool/proxy"
+	"github.com/zu1k/proxypool/tool"
 )
 
+var NeedFetchNewConfigFile = false
+
 func Crawl() {
+	if NeedFetchNewConfigFile {
+		FetchNewConfigFileThenInit()
+	}
 	proxies := make([]proxy.Proxy, 0)
 	for _, g := range Getters {
 		proxies = append(proxies, g.Get()...)
@@ -29,6 +37,9 @@ func Crawl() {
 }
 
 func CrawlGo() {
+	if NeedFetchNewConfigFile {
+		FetchNewConfigFileThenInit()
+	}
 	wg := &sync.WaitGroup{}
 	var pc = make(chan proxy.Proxy)
 	for _, g := range Getters {
@@ -55,3 +66,20 @@ func CrawlGo() {
 	cache.SetProxies(proxies)
 	cache.SetString("clashproxies", provider.Clash{Proxies: proxies}.Provide())
 }
+
+func FetchNewConfigFileThenInit() {
+	resp, err := tool.GetHttpClient().Get("https://raw.githubusercontent.com/zu1k/proxypool/master/source.yaml")
+	if err != nil {
+		return
+	}
+	defer resp.Body.Close()
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return
+	}
+	err = ioutil.WriteFile("source.yaml", body, os.ModePerm)
+	if err != nil {
+		return
+	}
+	InitConfigAndGetters("source.yaml")
+}
diff --git a/config/config.go b/config/config.go
index 15d717f..685d3ac 100644
--- a/config/config.go
+++ b/config/config.go
@@ -25,7 +25,7 @@ func Parse(path string) (*Config, error) {
 	if err != nil {
 		return nil, err
 	}
-
+	SourceConfig = Config{}
 	err = yaml.Unmarshal(fileData, &SourceConfig)
 	if err != nil {
 		return nil, err
diff --git a/main.go b/main.go
index c679d53..4294b63 100644
--- a/main.go
+++ b/main.go
@@ -3,27 +3,19 @@ package main
 import (
 	"flag"
 	"fmt"
-	"os"
 
 	"github.com/zu1k/proxypool/api"
 	"github.com/zu1k/proxypool/app"
-	"github.com/zu1k/proxypool/config"
 )
 
 func main() {
-	filePath := flag.String("c", "source.yaml", "path to config file: source.yaml")
+	filePath := flag.String("c", "", "path to config file: source.yaml")
 	flag.Parse()
-	c, err := config.Parse(*filePath)
-	if err != nil {
-		fmt.Println("Error: ", err.Error())
-		os.Exit(1)
+	if *filePath == "" {
+		app.NeedFetchNewConfigFile = true
+	} else {
+		app.InitConfigAndGetters(*filePath)
 	}
-	if c == nil {
-		fmt.Println("Error: no sources")
-		os.Exit(2)
-	}
-
-	app.InitGetters(c.Sources)
 	go app.Cron()
 	fmt.Println("Do the first crawl...")
 	app.CrawlGo()