Merge pull request #60 from Tribleave/master

Support proxy setting in config

cf.go

@@ -12,7 +12,7 @@ import (
 	docopt "github.com/docopt/docopt-go"
 )
 
-const version = "v0.8.2"
+const version = "v0.8.3"
 
 func main() {
 	usage := `Codeforces Tool $%version%$ (cf). https://github.com/xalanq/cf-tool

client/client.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 
@@ -21,6 +22,7 @@ type Client struct {
 	Bfaa           string          `json:"bfaa"`
 	LastSubmission *SaveSubmission `json:"last_submission"`
 	Host           string          `json:"host"`
+	Proxy          string          `json:"proxy"`
 	path           string
 	client         *http.Client
 }
@@ -39,15 +41,39 @@ func formatHost(host string) (string, error) {
 	return host, nil
 }
 
+func formatProxy(proxy string) (string, error) {
+	if len(proxy) == 0 {
+		return "", nil
+	}
+	reg := regexp.MustCompile(`(https?|socks5)://[\w\-\.]+:\d+`)
+	if !reg.MatchString(proxy) {
+		return "", fmt.Errorf(`Invalid proxy "%v"`, proxy)
+	}
+	for proxy[len(proxy)-1:] == "/" {
+		proxy = proxy[:len(proxy)-1]
+	}
+	return proxy, nil
+}
+
 // New client
 func New(path string) *Client {
 	jar, _ := cookiejar.New(nil)
 	c := &Client{Jar: jar, LastSubmission: nil, path: path, client: nil}
 	if path != "" {
 		c.load()
 	}
-	c.client = &http.Client{Jar: c.Jar}
-	var err error
+	proxyURL, err := formatProxy(c.Proxy)
+	proxy := http.ProxyFromEnvironment
+	if err != nil {
+		color.Red(err.Error() + `. Use default proxy from environment`)
+		color.Red(`Please use "cf config" to set a valid proxy later`)
+	} else if proxyURL != "" {
+		proxyURL, err := url.Parse(c.Proxy)
+		if err == nil {
+			proxy = http.ProxyURL(proxyURL)
+		}
+	}
+	c.client = &http.Client{Jar: c.Jar, Transport: &http.Transport{Proxy: proxy}}
 	c.Host, err = formatHost(c.Host)
 	if err != nil {
 		color.Red(err.Error() + `. Use default host "https://codeforces.com"`)
@@ -106,3 +132,33 @@ func (c *Client) SetHost() (err error) {
 	color.Green("New host domain is %v", host)
 	return c.save()
 }
+
+// SetProxy set proxy for client
+func (c *Client) SetProxy() (err error) {
+	proxy, err := formatProxy(c.Proxy)
+	if err != nil {
+		proxy = ""
+	}
+	if len(proxy) == 0 {
+		color.Green("Current proxy is based on environment")
+	} else {
+		color.Green("Current proxy is %v", proxy)
+	}
+	color.Cyan(`Set a new proxy (e.g. "http://127.0.0.1:80", "socks5://127.0.0.1:1080"`)
+	color.Cyan(`Enter empty line if you want to use default proxy from environment`)
+	color.Cyan(`Note: Proxy URL should match "proxyProtocol://proxyIp:proxyPort"`)
+	for {
+		proxy, err = formatProxy(util.ScanlineTrim())
+		if err == nil {
+			break
+		}
+		color.Red(err.Error())
+	}
+	c.Proxy = proxy
+	if len(proxy) == 0 {
+		color.Green("Current proxy is based on environment")
+	} else {
+		color.Green("Current proxy is %v", proxy)
+	}
+	return c.save()
+}

client/#.go

@@ -4,7 +4,6 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
-	"net/http"
 	"net/url"
 	"regexp"
 
@@ -49,7 +48,7 @@ func (c *Client) Login(username, password string) (err error) {
 	jar, _ := cookiejar.New(nil)
 	color.Cyan("Login %v...\n", username)
 
-	c.client = &http.Client{Jar: jar}
+	c.client.Jar = jar
 
 	resp, err := c.client.Get(c.Host + "/enter")
 	if err != nil {

client/parse.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"html"
 	"io/ioutil"
-	"net/http"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -45,8 +44,7 @@ func findSample(body []byte) (input [][]byte, output [][]byte, err error) {
 
 // ParseProblem parse problem to path
 func (c *Client) ParseProblem(URL, path string) (samples int, err error) {
-	client := &http.Client{Jar: c.Jar.Copy()}
-	resp, err := client.Get(URL)
+	resp, err := c.client.Get(URL)
 	if err != nil {
 		return
 	}

client/pull.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"html"
 	"io/ioutil"
-	"net/http"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -52,8 +51,7 @@ func (c *Client) PullCode(contestID, submissionID, path, ext string, rename bool
 	}
 
 	URL := ToGym(fmt.Sprintf(c.Host+"/contest/%v/submission/%v", contestID, submissionID), contestID)
-	client := &http.Client{Jar: c.Jar.Copy()}
-	resp, err := client.Get(URL)
+	resp, err := c.client.Get(URL)
 	if err != nil {
 		return
 	}

cmd/config.go

@@ -17,7 +17,8 @@ func Config(args map[string]interface{}) error {
 	ansi.Println(`3) set default template`)
 	ansi.Println(`4) run "cf gen" after "cf parse"`)
 	ansi.Println(`5) set host domain`)
-	index := util.ChooseIndex(6)
+	ansi.Println(`6) set proxy`)
+	index := util.ChooseIndex(7)
 	if index == 0 {
 		return config.New(config.ConfigPath).Login(config.SessionPath)
 	} else if index == 1 {
@@ -30,6 +31,8 @@ func Config(args map[string]interface{}) error {
 		return config.New(config.ConfigPath).SetGenAfterParse()
 	} else if index == 5 {
 		return client.New(config.SessionPath).SetHost()
+	} else if index == 6 {
+		return client.New(config.SessionPath).SetProxy()
 	}
 	return nil
 }