Skip to content

Commit

Permalink
feat(client): add uuid for more randomize token
Browse files Browse the repository at this point in the history
feat(client): add feature for upload more files
  • Loading branch information
ikr4-m committed Aug 8, 2024
1 parent a816ec6 commit ad500d0
Show file tree
Hide file tree
Showing 7 changed files with 183 additions and 123 deletions.
118 changes: 17 additions & 101 deletions cmd/mdrop-client/get_command.go
Original file line number Diff line number Diff line change
@@ -1,24 +1,19 @@
package main

import (
"bufio"
"crypto/sha256"
"errors"
"flag"
"fmt"
"io"
"net/http"
"os"
"strings"

"github.com/mplus-oss/mdrop/internal"
"github.com/schollz/progressbar/v3"
)

func GetCommand(args []string) {
errChan := make(chan error, 1)

reader := bufio.NewReader(os.Stdin)
flag := flag.NewFlagSet("mdrop get", flag.ExitOnError)
var (
help = flag.Bool("help", false, "Print this message")
Expand Down Expand Up @@ -70,108 +65,29 @@ func GetCommand(args []string) {
// Check tunnel
fmt.Println("Connecting to sender...")
GetTcpReadyConnect(*localPort)
client := http.Client{}
resp, err := client.Get(
fmt.Sprintf("http://localhost:%v/checksum", *localPort),
)
if err != nil {
internal.PrintErrorWithExit("sendHttpClientChecksum", err, 1)
}
if resp.StatusCode != http.StatusOK {
internal.PrintErrorWithExit("sendHttpClientResponseChecksum", err, 1)
}
checksumBytes, err := io.ReadAll(resp.Body)
if err != nil {
internal.PrintErrorWithExit("sendHttpClientReadChecksum", err, 1)
}
checksum := string(checksumBytes)

resp, err = client.Post(
fmt.Sprintf("http://localhost:%v/receive", *localPort),
"binary/octet-stream",
nil,
)
if err != nil {
internal.PrintErrorWithExit("sendHttpClient", err, 1)
}
if resp.StatusCode != http.StatusOK {
internal.PrintErrorWithExit("sendHttpClientResponse", err, 1)
}
defer resp.Body.Close()

// Set filename from header or from output
fileName := resp.Header.Get("X-Attachment-Name")
if fileName == "" {
internal.PrintErrorWithExit("sendHttpClientInvalidAttachmentName", err, 1)
}
if *fileNameOpt != "" {
fileName = *fileNameOpt
}
fmt.Println("File found:", fileName)

// Check if there's duplicate file
filePath, err := os.Getwd()
if err != nil {
internal.PrintErrorWithExit("sendFileWorkDir", err, 1)
}
if fileStatus, _ := os.Stat(filePath+"/"+fileName); fileStatus != nil {
fmt.Print("There's duplicate file. Action? [(R)eplace/R(e)name/(C)ancel] [Default: R] -> ")
prompt, err := reader.ReadString('\n')
// Downloading file
for _, uuid := range sender.Files {
// No error checking needed.
checksum := GetChecksum(*localPort, uuid)
filePath := GetDownload(*localPort, *fileNameOpt, uuid)

// Check checksum
fmt.Println("Checking checksum...")
fileDownloaded, err := os.Open(filePath)
if err != nil {
internal.PrintErrorWithExit("sendPromptError", err, 1)
internal.PrintErrorWithExit("checksumFileOpen", err, 1)
}
prompt = strings.Replace(prompt, "\n", "", -1)
if strings.ToLower(prompt) == "c" {
internal.PrintErrorWithExit("sendPromptCancel", errors.New("Canceled by action"), 0)
hash := sha256.New()
if _, err := io.Copy(hash, fileDownloaded); err != nil {
internal.PrintErrorWithExit("checksumHashSum", err, 1)
}
if strings.ToLower(prompt) == "e" {
fmt.Print("Change filename ["+fileName+"]: ")
prompt, err = reader.ReadString('\n')
if err != nil {
internal.PrintErrorWithExit("sendPromptError", err, 1)
}
prompt = strings.Replace(prompt, "\n", "", -1)
if prompt == fileName {
internal.PrintErrorWithExit("sendPromptDuplicateFilename", errors.New("Canceled by action"), 0)
}
fileName = prompt
checksumLocal := fmt.Sprintf("%x", hash.Sum(nil))
if checksumLocal != checksum {
internal.PrintErrorWithExit("checksumMismatch", errors.New("Checksum mismatch with sender"), 1)
}
fileDownloaded.Close()
}
filePath += "/"+fileName

// Create file
file, err := os.Create(filePath)
if err != nil {
internal.PrintErrorWithExit("sendFileCreation", err, 1)
}

// Downloading file
fmt.Println("Downloading...")
bar := progressbar.DefaultBytes(resp.ContentLength, fileName)
_, err = io.Copy(io.MultiWriter(bar, file), resp.Body)
if err != nil {
errMsg := err.Error()
if strings.Contains(errMsg, "EOF") {
err = errors.New("Broken pipe from sender because forced close or terminated.")
}
internal.PrintErrorWithExit("sendStreamFile", err, 1)
}

// Check checksum
fmt.Println("Checking checksum...")
fileDownloaded, err := os.Open(filePath)
if err != nil {
internal.PrintErrorWithExit("checksumFileOpen", err, 1)
}
hash := sha256.New()
if _, err := io.Copy(hash, fileDownloaded); err != nil {
internal.PrintErrorWithExit("checksumHashSum", err, 1)
}
checksumLocal := fmt.Sprintf("%x", hash.Sum(nil))
if checksumLocal != checksum {
internal.PrintErrorWithExit("checksumMismatch", errors.New("Checksum mismatch with sender"), 1)
}
fileDownloaded.Close()

fmt.Println("Download success!")
}
114 changes: 114 additions & 0 deletions cmd/mdrop-client/get_download.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
package main

import (
"bufio"
"errors"
"fmt"
"io"
"net/http"
"os"
"strings"

"github.com/mplus-oss/mdrop/internal"
"github.com/schollz/progressbar/v3"
)

func GetChecksum(localPort int, uuid string) string {
client := http.Client{}
resp, err := client.Get(
fmt.Sprintf("http://localhost:%v/checksum-%v", localPort, uuid),
)
if err != nil {
internal.PrintErrorWithExit("sendHttpClientChecksum", err, 1)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
err = internal.CustomizeError("sendHttpClientResponseChecksum", errors.New("Checksum response error"))
internal.PrintErrorWithExit("sendHttpClientResponseChecksum", err, 1)
}
checksumBytes, err := io.ReadAll(resp.Body)
if err != nil {
internal.PrintErrorWithExit("sendHttpClientReadChecksum", err, 1)
}
return string(checksumBytes)
}

func GetDownload(localPort int, fileNameOpt string, uuid string) string {
reader := bufio.NewReader(os.Stdin)
client := http.Client{}

resp, err := client.Post(
fmt.Sprintf("http://localhost:%v/%v", localPort, uuid),
"binary/octet-stream",
nil,
)
if err != nil {
internal.PrintErrorWithExit("sendHttpClient", err, 1)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
fmt.Println("Status Code:", resp.StatusCode)
err = internal.CustomizeError("sendHttpClientResponse", errors.New("Download response error"))
internal.PrintErrorWithExit("sendHttpClientResponse", err, 1)
}

// Set filename from header or from output
fileName := resp.Header.Get("X-Attachment-Name")
if fileName == "" {
internal.PrintErrorWithExit("sendHttpClientInvalidAttachmentName", err, 1)
}
if fileNameOpt != "" {
fileName = fileNameOpt
}
fmt.Println("File found:", fileName)

// Check if there's duplicate file
filePath, err := os.Getwd()
if err != nil {
internal.PrintErrorWithExit("sendFileWorkDir", err, 1)
}
if fileStatus, _ := os.Stat(filePath+"/"+fileName); fileStatus != nil {
fmt.Print("There's duplicate file. Action? [(R)eplace/R(e)name/(C)ancel] [Default: R] -> ")
prompt, err := reader.ReadString('\n')
if err != nil {
internal.PrintErrorWithExit("sendPromptError", err, 1)
}
prompt = strings.Replace(prompt, "\n", "", -1)
if strings.ToLower(prompt) == "c" {
internal.PrintErrorWithExit("sendPromptCancel", errors.New("Canceled by action"), 0)
}
if strings.ToLower(prompt) == "e" {
fmt.Print("Change filename ["+fileName+"]: ")
prompt, err = reader.ReadString('\n')
if err != nil {
internal.PrintErrorWithExit("sendPromptError", err, 1)
}
prompt = strings.Replace(prompt, "\n", "", -1)
if prompt == fileName {
internal.PrintErrorWithExit("sendPromptDuplicateFilename", errors.New("Canceled by action"), 0)
}
fileName = prompt
}
}
filePath += "/"+fileName

// Create file
file, err := os.Create(filePath)
if err != nil {
internal.PrintErrorWithExit("sendFileCreation", err, 1)
}

// Downloading file
fmt.Println("Downloading...")
bar := progressbar.DefaultBytes(resp.ContentLength, fileName)
_, err = io.Copy(io.MultiWriter(bar, file), resp.Body)
if err != nil {
errMsg := err.Error()
if strings.Contains(errMsg, "EOF") {
err = errors.New("Broken pipe from sender because forced close or terminated.")
}
internal.PrintErrorWithExit("sendStreamFile", err, 1)
}

return filePath
}
36 changes: 24 additions & 12 deletions cmd/mdrop-client/send_command.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"strconv"
"strings"

"github.com/google/uuid"
"github.com/mplus-oss/mdrop/internal"
)

Expand All @@ -23,8 +24,8 @@ func SendCommand(args []string) {
)
flag.Parse(args)

file := flag.Arg(0)
if *help || file == "" {
file := flag.Args()
if *help || len(file) == 0 {
fmt.Println("Command: mdrop send [options] <file>")
flag.Usage()
os.Exit(1)
Expand Down Expand Up @@ -76,21 +77,32 @@ func SendCommand(args []string) {
}
}()

// Deploy webserver instance
// Deploy webserver instance & add uuid for resolver
go func() {
fileUUID := []string{}

for range file {
fileUUID = append(fileUUID, uuid.New().String())
}

// Print token
token, err := TokenTransferJSON{
Host: config.Host,
RemotePort: remotePort,
Files: fileUUID,
}.GenerateToken()
if err != nil {
errChan <- err
}

fmt.Print("\nPlease copy this token to the receiver.")
fmt.Print("\nToken: "+token+"\n\n")

fmt.Println("Spawning webserver...")
err := SendWebserver(*localPort, file)
err = SendWebserver(*localPort, file, fileUUID)
errChan <- err
}()

// Print token
token, err := TokenTransferJSON{Host: config.Host, RemotePort: remotePort}.GenerateToken()
if err != nil {
internal.PrintErrorWithExit("generateTokenError", err, 1)
}
fmt.Print("\nPlease copy this token to the receiver.")
fmt.Print("\nToken: "+token+"\n\n")

// Check if there's some error on different threads
err = <-errChan
if err != nil {
Expand Down
30 changes: 22 additions & 8 deletions cmd/mdrop-client/send_webserver.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,24 @@ import (
)

var server *http.Server = &http.Server{}
var filePath string = ""
var isStillUsed bool = false
var totalFile int = 0
var fileDownloaded int = 0

var senderErrorChan chan error = make(chan error)

func SendWebserver(localPort int, file string) (err error) {
filePath = file
func SendWebserver(localPort int, file []string, uuid []string) (err error) {
totalFile = len(file)
server.Addr = ":"+strconv.Itoa(localPort)

http.Handle("/receive", http.HandlerFunc(receiveSendWebserver))
http.Handle("/checksum", http.HandlerFunc(checksumSendWebserver))
for i, _ := range file {
http.Handle("/"+uuid[i], http.HandlerFunc(func (w http.ResponseWriter, request *http.Request) {
receiveSendWebserver(w, request, file[i])
}))
http.Handle("/checksum-"+uuid[i], http.HandlerFunc(func (w http.ResponseWriter, request *http.Request) {
checksumSendWebserver(w, request, file[i])
}))
}

go func() {
err := server.ListenAndServe()
Expand All @@ -52,7 +59,7 @@ func SendWebserver(localPort int, file string) (err error) {
return nil
}

func checksumSendWebserver(w http.ResponseWriter, request *http.Request) {
func checksumSendWebserver(w http.ResponseWriter, request *http.Request, filePath string) {
fmt.Println("Receiver taking the checksum file.")
file, err := os.Open(filePath)
if err != nil {
Expand All @@ -72,7 +79,7 @@ func checksumSendWebserver(w http.ResponseWriter, request *http.Request) {
request.Close = true
}

func receiveSendWebserver(w http.ResponseWriter, request *http.Request) {
func receiveSendWebserver(w http.ResponseWriter, request *http.Request, filePath string) {
if request.Method != "POST" {
w.WriteHeader(http.StatusMethodNotAllowed)
return
Expand Down Expand Up @@ -119,5 +126,12 @@ func receiveSendWebserver(w http.ResponseWriter, request *http.Request) {
}

request.Close = true
senderErrorChan <- nil

// Send channel to shutdown if file downloaded is same as file sent
fileDownloaded += 1
if fileDownloaded == totalFile {
senderErrorChan <- nil
} else {
isStillUsed = false
}
}
Loading

0 comments on commit ad500d0

Please sign in to comment.