bitbucket.org/Aishee/synsec@v0.0.0-20210414005726-236fc01a153d/pkg/types/dataset.go (about)

     1  package types
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"net/http"
     7  	"os"
     8  	"path"
     9  
    10  	log "github.com/sirupsen/logrus"
    11  )
    12  
    13  type DataSource struct {
    14  	SourceURL string `yaml:"source_url"`
    15  	DestPath  string `yaml:"dest_file"`
    16  	Type      string `yaml:"type"`
    17  }
    18  
    19  type DataSet struct {
    20  	Data []*DataSource `yaml:"data,omitempty"`
    21  }
    22  
    23  func downloadFile(url string, destPath string) error {
    24  	log.Debugf("downloading %s in %s", url, destPath)
    25  	req, err := http.NewRequest("GET", url, nil)
    26  	if err != nil {
    27  		return err
    28  	}
    29  
    30  	resp, err := http.DefaultClient.Do(req)
    31  	if err != nil {
    32  		return err
    33  	}
    34  	defer resp.Body.Close()
    35  
    36  	body, err := ioutil.ReadAll(resp.Body)
    37  	if err != nil {
    38  		return err
    39  	}
    40  
    41  	if resp.StatusCode != 200 {
    42  		return fmt.Errorf("download response 'HTTP %d' : %s", resp.StatusCode, string(body))
    43  	}
    44  
    45  	file, err := os.OpenFile(destPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
    46  	if err != nil {
    47  		return err
    48  	}
    49  
    50  	_, err = file.WriteString(string(body))
    51  	if err != nil {
    52  		return err
    53  	}
    54  
    55  	err = file.Sync()
    56  	if err != nil {
    57  		return err
    58  	}
    59  
    60  	return nil
    61  }
    62  
    63  func GetData(data []*DataSource, dataDir string) error {
    64  	for _, dataS := range data {
    65  		destPath := path.Join(dataDir, dataS.DestPath)
    66  		log.Infof("downloading data '%s' in '%s'", dataS.SourceURL, destPath)
    67  		err := downloadFile(dataS.SourceURL, destPath)
    68  		if err != nil {
    69  			return err
    70  		}
    71  	}
    72  
    73  	return nil
    74  }