From 210fb1e02453413d1ce070b70c850807286a1a7a Mon Sep 17 00:00:00 2001
From: jet2tlf <jet2tlf@gmail.com>
Date: Mon, 3 Jun 2024 15:31:42 -0300
Subject: codecrafters submit [skip ci]

---
 cmd/mybittorrent/main.go | 32 ++++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

(limited to 'cmd/mybittorrent/main.go')

diff --git a/cmd/mybittorrent/main.go b/cmd/mybittorrent/main.go
index 80a9b7d..6dbdda6 100644
--- a/cmd/mybittorrent/main.go
+++ b/cmd/mybittorrent/main.go
@@ -75,6 +75,9 @@ func main() {
 		peerAddr := os.Args[3]
 
 		peer, err := c.Handshake(fn, peerAddr)
+
+		defer peer.Close()
+
 		if err != nil {
 			panic(err)
 		}
@@ -91,6 +94,7 @@ func main() {
 		}
 
 		c := createClient(fn)
+		defer c.Close()
 		pr, err := c.GetPeers(fn)
 		if err != nil {
 			panic(err)
@@ -109,15 +113,39 @@ func main() {
 			panic(fmt.Errorf("no peers found for file: %s", out))
 		}
 
-		defer peer.Close()
+		f, err := os.Create(out)
+		if err != nil {
+			panic(f)
+		}
 
-		err = peer.DownloadPiece(out, index)
+		defer f.Close()
+
+		err = peer.DownloadPiece(f, c.Torrents[fn].Meta.Pieces()[index])
 		if err != nil {
 			panic(err)
 		}
 
 		fmt.Printf("Piece %d downloaded to %s.", index, out)
 
+	case "download":
+		out := os.Args[3]
+		fn := os.Args[4]
+		c := createClient(fn)
+
+		defer c.Close()
+
+		ct, err := c.ConnectPeers(fn)
+		if err != nil {
+			panic(err)
+		}
+
+		err = ct.Download(out)
+		if err != nil {
+			panic(err)
+		}
+
+		fmt.Printf("Downloaded %s to %s", fn, out)
+
 	default:
 		fmt.Println("Unknown command: " + command)
 		os.Exit(1)
-- 
cgit v1.2.3