jetz
4/21/2018 - 12:42 PM

File transfer between server and local

File transfer between server and local

package main

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/urfave/cli"
	"gopkg.in/cheggaaa/pb.v1"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"strings"
)

const (
	// PORT server port
	PORT = 17621
	// ADDR remote server addr
	ADDR = "127.0.0.1"
	// TOKEN request auth token
	TOKEN = "urToken"
)

func filesHandler(w http.ResponseWriter, r *http.Request) {
	for _, token := range r.Header["X-Ft-Token"] {
		if token == TOKEN {
			switch r.Method {
			case "POST":
				err := r.ParseMultipartForm(32 << 20)
				if err != nil {
					http.Error(w, err.Error(), http.StatusInternalServerError)
					return
				}

				files := r.MultipartForm.File["uploadfile"]

				for _, file := range files {
					src, err := file.Open()
					if err != nil {
						http.Error(w, err.Error(), http.StatusInternalServerError)
						return
					}
					defer src.Close()

					dst, err := os.Create("/data/ft/" + file.Filename)
					if err != nil {
						http.Error(w, err.Error(), http.StatusInternalServerError)
						return
					}
					defer dst.Close()

					if _, err := io.Copy(dst, src); err != nil {
						http.Error(w, err.Error(), http.StatusInternalServerError)
						return
					}
				}
			case "GET":
				fs := http.FileServer(http.Dir("/data/ft"))
				http.StripPrefix("/files/", fs).ServeHTTP(w, r)
			case "DELETE":
				paths := strings.Split(r.URL.Path, "/")
				file := paths[len(paths)-1]
				err := os.Remove("/data/ft/" + file)
				if err != nil {
					http.Error(w, err.Error(), http.StatusInternalServerError)
					return
				}
			case "HEAD":
				files, err := ioutil.ReadDir("/data/ft")
				if err != nil {
					http.Error(w, err.Error(), http.StatusInternalServerError)
					return
				}
				var fileNames []string
				for _, f := range files {
					fileNames = append(fileNames, f.Name())
				}
				w.Header().Set("X-File-List", strings.Join(fileNames, "/"))
				w.Write(nil)
			default:
				w.WriteHeader(http.StatusMethodNotAllowed)
				w.Write([]byte("Unsupport request method"))
			}
			return
		}
	}
	w.WriteHeader(http.StatusForbidden)
	w.Write([]byte("Wrong token: forbidden"))
}

// Server runserver
type Server struct {
	Addr string
	Port uint
}

func (s *Server) start() {
	fmt.Println(fmt.Sprintf("FT Server running @%s:%d...", s.Addr, s.Port))
	http.HandleFunc("/files/", filesHandler)
	http.ListenAndServe(fmt.Sprintf("%s:%d", s.Addr, s.Port), nil)
}

// Client send/recv/list/remove
type Client struct {
	Addr string
	Port uint
}

func (c *Client) send(files []string) {
	uploadFile := func(url, file string) error {
		fr, err := os.Open(file)
		if err != nil {
			return err
		}
		defer fr.Close()

		buf := &bytes.Buffer{}
		mw := multipart.NewWriter(buf)
		fw, err := mw.CreateFormFile("uploadfile", file)
		if err != nil {
			return err
		}

		_, err = io.Copy(fw, fr)
		if err != nil {
			return err
		}

		mw.Close()

		fileStat, _ := fr.Stat()
		fileSize := fileStat.Size()
		bar := pb.New(int(fileSize)).SetUnits(pb.U_BYTES).Prefix(file)
		defer bar.Finish()
		bar.ShowSpeed = true
		bar.ShowTimeLeft = true
		bar.Start()
		src := bar.NewProxyReader(buf)

		req, err := http.NewRequest("POST", url, src)
		if err != nil {
			return err
		}
		req.Header.Set("X-Ft-Token", TOKEN)
		req.Header.Set("Content-Type", mw.FormDataContentType())

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}

		if resp.StatusCode != http.StatusOK {
			return errors.New("Bad status: " + resp.Status)
		}
		return nil
	}

	url := fmt.Sprintf("http://%s:%d/files/", c.Addr, c.Port)
	for _, file := range files {
		if err := uploadFile(url, file); err != nil {
			fmt.Println(err)
		}
	}
}

func (c *Client) recv(files []string) {
	downloadFile := func(url, file string) error {
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			return err
		}
		req.Header.Set("X-Ft-Token", TOKEN)

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		if resp.StatusCode == http.StatusOK {
			out, err := os.Create(file)
			if err != nil {
				return err
			}
			defer out.Close()

			bar := pb.New(int(resp.ContentLength)).SetUnits(pb.U_BYTES).Prefix(file)
			defer bar.Finish()
			bar.ShowSpeed = true
			bar.ShowTimeLeft = true
			bar.Start()

			src := bar.NewProxyReader(resp.Body)
			_, err = io.Copy(out, src)
			if err != nil {
				return err
			}
		} else if resp.StatusCode == http.StatusNotFound {
			return errors.New(file + " Not Found")
		} else {
			return errors.New("Bad status: " + resp.Status)
		}
		return nil
	}

	for _, file := range files {
		url := fmt.Sprintf("http://%s:%d/files/%s", c.Addr, c.Port, file)
		if err := downloadFile(url, file); err != nil {
			fmt.Println(err)
		}
	}
}

func (c *Client) remove(files []string) {
	removeFile := func(url, file string) error {
		req, err := http.NewRequest("DELETE", url, nil)
		if err != nil {
			return err
		}
		req.Header.Set("X-Ft-Token", TOKEN)

		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusNotFound {
			return errors.New(file + " Not Found")
		}
		return nil
	}
	for _, file := range files {
		url := fmt.Sprintf("http://%s:%d/files/%s", c.Addr, c.Port, file)
		if err := removeFile(url, file); err != nil {
			fmt.Println(err)
		}
	}
}

func (c *Client) list() {
	url := fmt.Sprintf("http://%s:%d/files/", c.Addr, c.Port)
	req, err := http.NewRequest("HEAD", url, nil)
	if err != nil {
		return
	}
	req.Header.Set("X-Ft-Token", TOKEN)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusOK {
		fileStr := resp.Header.Get("X-File-List")
		if len(fileStr) != 0 {
			files := strings.Split(fileStr, "/")
			fmt.Println(len(files))
			fmt.Println("=============================")
			for _, file := range files {
				fmt.Println(file)
			}
			fmt.Println("=============================")
		}
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "ft"
	app.Usage = "file transfer between server and local"
	app.Version = "1.0.0"
	app.Commands = []cli.Command{
		{
			Name:  "runserver",
			Usage: "file transfer server mode running",
			Action: func(c *cli.Context) error {
				(&Server{"0.0.0.0", PORT}).start()
				return nil
			},
		},
		{
			Name:    "send",
			Aliases: []string{"sz"},
			Usage:   "send files to server",
			Action: func(c *cli.Context) error {
				(&Client{ADDR, PORT}).send(c.Args())
				return nil
			},
		},
		{
			Name:    "recv",
			Aliases: []string{"rz"},
			Usage:   "receive files from server",
			Action: func(c *cli.Context) error {
				(&Client{ADDR, PORT}).recv(c.Args())
				return nil
			},
		},
		{
			Name:    "remove",
			Aliases: []string{"rm"},
			Usage:   "rm files stored in server",
			Action: func(c *cli.Context) error {
				(&Client{ADDR, PORT}).remove(c.Args())
				return nil
			},
		},
		{
			Name:    "list",
			Aliases: []string{"ls"},
			Usage:   "list files stored in server",
			Action: func(c *cli.Context) error {
				(&Client{ADDR, PORT}).list()
				return nil
			},
		},
	}
	app.Run(os.Args)
}