1package fetcher
2
3import (
4	"bytes"
5	"compress/bzip2"
6	"fmt"
7	"io"
8	"net/http"
9	"net/url"
10	"sync"
11	"time"
12
13	"github.com/htcat/htcat"
14	"github.com/inconshreveable/log15"
15	c "github.com/kotakanbe/goval-dictionary/config"
16	"github.com/kotakanbe/goval-dictionary/util"
17)
18
19type fetchRequest struct {
20	target       string
21	url          string
22	bzip2        bool
23	concurrently bool
24}
25
26//FetchResult has url and OVAL definitions
27type FetchResult struct {
28	Target string
29	URL    string
30	Body   []byte
31}
32
33func fetchFeedFiles(reqs []fetchRequest) (results []FetchResult, err error) {
34	reqChan := make(chan fetchRequest, len(reqs))
35	resChan := make(chan FetchResult, len(reqs))
36	errChan := make(chan error, len(reqs))
37	defer close(reqChan)
38	defer close(resChan)
39	defer close(errChan)
40
41	for _, r := range reqs {
42		log15.Info("Fetching... ", "URL", r.url)
43	}
44
45	go func() {
46		for _, r := range reqs {
47			reqChan <- r
48		}
49	}()
50
51	concurrency := len(reqs)
52	tasks := util.GenWorkers(concurrency)
53	wg := new(sync.WaitGroup)
54	for range reqs {
55		wg.Add(1)
56		tasks <- func() {
57			select {
58			case req := <-reqChan:
59				var body []byte
60				if req.concurrently {
61					body, err = fetchFileConcurrently(req, 20/len(reqs))
62				} else {
63					body, err = fetchFileWithUA(req)
64				}
65				wg.Done()
66				if err != nil {
67					errChan <- err
68					return
69				}
70				resChan <- FetchResult{
71					Target: req.target,
72					URL:    req.url,
73					Body:   body,
74				}
75			}
76			return
77		}
78	}
79	wg.Wait()
80
81	errs := []error{}
82	timeout := time.After(10 * 60 * time.Second)
83	for range reqs {
84		select {
85		case res := <-resChan:
86			results = append(results, res)
87			log15.Info("Fetched... ", "URL", res.URL)
88		case err := <-errChan:
89			errs = append(errs, err)
90		case <-timeout:
91			return results, fmt.Errorf("Timeout Fetching")
92		}
93	}
94	log15.Info("Finished fetching OVAL definitions")
95	if 0 < len(errs) {
96		return results, fmt.Errorf("%s", errs)
97	}
98	return results, nil
99}
100
101func fetchFileConcurrently(req fetchRequest, concurrency int) (body []byte, err error) {
102	var proxyURL *url.URL
103	httpCilent := &http.Client{}
104	if c.Conf.HTTPProxy != "" {
105		if proxyURL, err = url.Parse(c.Conf.HTTPProxy); err != nil {
106			return nil, fmt.Errorf("Failed to parse proxy url: %s", err)
107		}
108		httpCilent = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}}
109	}
110
111	u, err := url.Parse(req.url)
112	if err != nil {
113		return nil, fmt.Errorf("aborting: could not parse given URL: %v", err)
114	}
115
116	buf := bytes.Buffer{}
117	htc := htcat.New(httpCilent, u, concurrency)
118	if _, err := htc.WriteTo(&buf); err != nil {
119		return nil, fmt.Errorf("aborting: could not write to output stream: %v",
120			err)
121	}
122
123	var bytesBody []byte
124	if req.bzip2 {
125		var b bytes.Buffer
126		if _, err := b.ReadFrom(bzip2.NewReader(bytes.NewReader(buf.Bytes()))); err != nil {
127			return body, err
128		}
129		bytesBody = b.Bytes()
130	} else {
131		bytesBody = buf.Bytes()
132	}
133
134	return bytesBody, nil
135}
136
137func fetchFileWithUA(req fetchRequest) (body []byte, err error) {
138	var errs []error
139	var proxyURL *url.URL
140	var resp *http.Response
141
142	httpClient := &http.Client{}
143	if c.Conf.HTTPProxy != "" {
144		if proxyURL, err = url.Parse(c.Conf.HTTPProxy); err != nil {
145			return nil, fmt.Errorf("Failed to parse proxy url: %s", err)
146		}
147		httpClient = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}}
148	}
149
150	httpreq, err := http.NewRequest("GET", req.url, nil)
151	if err != nil {
152		return nil, fmt.Errorf("Download failed: %v", err)
153	}
154
155	httpreq.Header.Set("User-Agent", "curl/7.37.0")
156	resp, err = httpClient.Do(httpreq)
157	if err != nil {
158		return nil, fmt.Errorf("Download failed: %v", err)
159	}
160	defer resp.Body.Close()
161
162	buf := bytes.NewBuffer(nil)
163	if _, err := io.Copy(buf, resp.Body); err != nil {
164		return nil, err
165	}
166	if len(errs) > 0 || resp == nil || resp.StatusCode != 200 {
167		return nil, fmt.Errorf(
168			"HTTP error. errs: %v, url: %s", errs, req.url)
169	}
170
171	var bytesBody []byte
172	if req.bzip2 {
173		bz := bzip2.NewReader(buf)
174		var b bytes.Buffer
175		if _, err := b.ReadFrom(bz); err != nil {
176			return nil, err
177		}
178		bytesBody = b.Bytes()
179	} else {
180		bytesBody = buf.Bytes()
181	}
182
183	return bytesBody, nil
184}
185