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