From b8b403f83eb019ffccaa51675d3ff5a44be23e07 Mon Sep 17 00:00:00 2001
From: Alexander Baranov <scr4t@yandex.ru>
Date: Mon, 11 May 2015 17:50:59 +0300
Subject: [PATCH] Basic pick-file functionality

---
 client.go            |  83 ++++++++++++++++++++++
 cmd/picker/picker.go | 160 +++++++++++++++++++++++++++++++++++++++++++
 torrent.go           |   9 +++
 3 files changed, 252 insertions(+)
 create mode 100644 cmd/picker/picker.go

diff --git a/client.go b/client.go
index f5ecac7c..9055955f 100644
--- a/client.go
+++ b/client.go
@@ -23,6 +23,7 @@ import (
 	"strings"
 	"syscall"
 	"time"
+	"math"
 
 	"github.com/anacrolix/sync"
 	"github.com/anacrolix/utp"
@@ -2016,6 +2017,62 @@ func (t Torrent) DownloadAll() {
 	t.cl.raisePiecePriority(t.torrent, t.numPieces()-1, piecePriorityReadahead)
 }
 
+// Marks the entire torrent for download. Requires the info first, see
+// GotInfo.
+func (t Torrent) DownloadFile(Path []string) {
+	t.cl.mu.Lock()
+	defer t.cl.mu.Unlock()
+
+	// log.Printf("File to Download: %s", Path)
+
+	// log.Printf("Pieces: %s", t.torrent.Info.NumPieces())
+	// log.Printf("Length: %s", t.torrent.Info.TotalLength())
+	// log.Printf("Torrent info: %s", t.torrent.Info.UpvertedFiles())
+
+	var offset int64
+	var pickedFile metainfo.FileInfo
+
+	found := false
+	pathStr := strings.Join(Path, "")
+
+	for _, file := range t.torrent.Info.UpvertedFiles() {
+		if strings.Join(file.Path, "/") ==  pathStr {
+			log.Printf("Found file: %s", file)
+
+			found = true
+			pickedFile = file
+			break
+		}
+		// log.Printf("%d %d `%s` `%s`", len(file.Path), len(Path), strings.Join(file.Path, "/"), strings.Join(Path, ""))
+		log.Printf("File: %s", strings.Join(file.Path, "/"))
+		offset += file.Length;
+	}
+
+	if !found {
+		panic(fmt.Sprintf("File not found"))
+	}
+
+	log.Printf("Donwloading file: `%s`", Path)
+	log.Printf("Calculated offset: %s", offset)
+	log.Printf("File length: %s", pickedFile.Length)
+	log.Printf("Piece length: %s", t.torrent.Info.PieceLength)
+
+
+	firstChunk := int(offset/t.torrent.Info.PieceLength)
+	nChunks := int(math.Ceil(float64(pickedFile.Length) / float64(t.torrent.Info.PieceLength)))
+
+	log.Printf("First chunk: %s", offset/t.torrent.Info.PieceLength)
+	log.Printf("Number of chunks: %s", nChunks)
+	log.Printf("Total chunks: %s", t.torrent.Info.NumPieces())
+
+
+	for chunk := firstChunk; chunk < firstChunk + nChunks; chunk++ {
+		log.Printf("Piece #%d: %s %s", chunk, t.torrent.Info.Piece(chunk).Length(), t.torrent.Info.Piece(chunk).Offset())
+		t.cl.raisePiecePriority(t.torrent, chunk, piecePriorityNormal)
+	}
+}
+
+
 // Returns nil metainfo if it isn't in the cache. Checks that the retrieved
 // metainfo has the correct infohash.
 func (cl *Client) torrentCacheMetaInfo(ih InfoHash) (mi *metainfo.MetaInfo, err error) {
@@ -2400,6 +2457,18 @@ func (cl *Client) allTorrentsCompleted() bool {
 	return true
 }
 
+func (cl *Client) allNeededTorrentsCompleted() bool {
+	for _, t := range cl.torrents {
+		if !t.haveInfo() {
+			return false
+		}
+		if ! t.neededPiecesDownloaded() {
+			return false
+		}
+	}
+	return true
+}
+
 // Returns true when all torrents are completely downloaded and false if the
 // client is stopped before that.
 func (me *Client) WaitAll() bool {
@@ -2414,6 +2483,20 @@ func (me *Client) WaitAll() bool {
 	return true
 }
 
+// Returns true when all requested chunks are completely downloaded and false if the
+// client is stopped before that.
+func (me *Client) WaitNeeded() bool {
+	me.mu.Lock()
+	defer me.mu.Unlock()
+	for !me.allNeededTorrentsCompleted() {
+		if me.stopped() {
+			return false
+		}
+		me.event.Wait()
+	}
+	return true
+}
+
 func (me *Client) fillRequests(t *torrent, c *connection) {
 	if c.Interested {
 		if c.PeerChoked {
diff --git a/cmd/picker/picker.go b/cmd/picker/picker.go
new file mode 100644
index 00000000..e882b62b
--- /dev/null
+++ b/cmd/picker/picker.go
@@ -0,0 +1,160 @@
+// Downloads torrents from the command-line.
+package main
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	_ "net/http/pprof"
+	"os"
+	"strings"
+	"time"
+
+	_ "github.com/anacrolix/envpprof"
+	"github.com/dustin/go-humanize"
+	"github.com/jessevdk/go-flags"
+
+	"github.com/anacrolix/torrent"
+	"github.com/anacrolix/torrent/metainfo"
+)
+
+// fmt.Fprintf(os.Stderr, "Usage: %s \n", os.Args[0])
+
+func resolvedPeerAddrs(ss []string) (ret []torrent.Peer, err error) {
+	for _, s := range ss {
+		var addr *net.TCPAddr
+		addr, err = net.ResolveTCPAddr("tcp", s)
+		if err != nil {
+			return
+		}
+		ret = append(ret, torrent.Peer{
+			IP:   addr.IP,
+			Port: addr.Port,
+		})
+	}
+	return
+}
+
+func bytesCompleted(tc *torrent.Client) (ret int64) {
+	for _, t := range tc.Torrents() {
+		if t.Info != nil {
+			ret += t.BytesCompleted()
+		}
+	}
+	return
+}
+
+// Returns an estimate of the total bytes for all torrents.
+func totalBytesEstimate(tc *torrent.Client) (ret int64) {
+	var noInfo, hadInfo int64
+	for _, t := range tc.Torrents() {
+		info := t.Info()
+		if info == nil {
+			noInfo++
+			continue
+		}
+		ret += info.TotalLength()
+		hadInfo++
+	}
+	if hadInfo != 0 {
+		// Treat each torrent without info as the average of those with,
+		// rounded up.
+		ret += (noInfo*ret + hadInfo - 1) / hadInfo
+	}
+	return
+}
+
+func progressLine(tc *torrent.Client) string {
+	return fmt.Sprintf("\033[K%s / %s\r", humanize.Bytes(uint64(bytesCompleted(tc))), humanize.Bytes(uint64(totalBytesEstimate(tc))))
+}
+
+func main() {
+	log.SetFlags(log.LstdFlags | log.Lshortfile)
+	var rootGroup struct {
+		Client    torrent.Config `group:"Client Options"`
+		Seed      bool           `long:"seed" description:"continue seeding torrents after completed"`
+		TestPeers []string       `long:"test-peer" description:"address of peer to inject to every torrent"`
+		Pick      []string       `long:"pick" description:"filename to pick"`
+	}
+	// Don't pass flags.PrintError because it's inconsistent with printing.
+	// https://github.com/jessevdk/go-flags/issues/132
+	parser := flags.NewParser(&rootGroup, flags.HelpFlag|flags.PassDoubleDash)
+	parser.Usage = "[OPTIONS] (magnet URI or .torrent file path)..."
+	posArgs, err := parser.Parse()
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "Download from the BitTorrent network.\n")
+		fmt.Println(err)
+		os.Exit(2)
+	}
+	log.Printf("File to pick: %s", rootGroup.Pick)
+
+	testPeers, err := resolvedPeerAddrs(rootGroup.TestPeers)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	if len(posArgs) == 0 {
+		fmt.Fprintln(os.Stderr, "no torrents specified")
+		return
+	}
+	client, err := torrent.NewClient(&rootGroup.Client)
+	if err != nil {
+		log.Fatalf("error creating client: %s", err)
+	}
+	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
+		client.WriteStatus(w)
+	})
+	defer client.Close()
+	for _, arg := range posArgs {
+		t := func() torrent.Torrent {
+			if strings.HasPrefix(arg, "magnet:") {
+				t, err := client.AddMagnet(arg)
+				if err != nil {
+					log.Fatalf("error adding magnet: %s", err)
+				}
+				return t
+			} else {
+				metaInfo, err := metainfo.LoadFromFile(arg)
+				if err != nil {
+					log.Fatal(err)
+				}
+				t, err := client.AddTorrent(metaInfo)
+				if err != nil {
+					log.Fatal(err)
+				}
+				return t
+			}
+		}()
+		err := t.AddPeers(testPeers)
+		if err != nil {
+			log.Fatal(err)
+		}
+		go func() {
+			<-t.GotInfo()
+			t.DownloadFile(rootGroup.Pick)
+		}()
+	}
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		if client.WaitNeeded() {
+			log.Print("downloaded ALL the torrents")
+		} else {
+			log.Fatal("y u no complete torrents?!")
+		}
+	}()
+	ticker := time.NewTicker(time.Second)
+waitDone:
+	for {
+		select {
+		case <-done:
+			break waitDone
+		case <-ticker.C:
+			os.Stdout.WriteString(progressLine(client))
+		}
+	}
+	if rootGroup.Seed {
+		select {}
+	}
+}
diff --git a/torrent.go b/torrent.go
index fd213d5b..d025185b 100644
--- a/torrent.go
+++ b/torrent.go
@@ -515,6 +515,15 @@ func (t *torrent) numPieces() int {
 	return t.Info.NumPieces()
 }
 
+func (t *torrent) neededPiecesDownloaded() bool {
+	for i := range iter.N(t.Info.NumPieces()) {
+		if t.Pieces[i].Priority != piecePriorityNone && !t.pieceComplete(i) {
+			return false
+		}
+	}
+	return true
+}
+
 func (t *torrent) numPiecesCompleted() (num int) {
 	for i := range iter.N(t.Info.NumPieces()) {
 		if t.pieceComplete(i) {
-- 
2.51.0