cmd/torrent-pick/main.go

   1 // Downloads torrents from the command-line.
   2 package main
   3
   4 import (
   5         "bufio"
   6         "fmt"
   7         "io"
   8         "log"
   9         "net"
  10         "net/http"
  11         _ "net/http/pprof"
  12         "os"
  13         "strings"
  14         "time"
  15
  16         _ "github.com/anacrolix/envpprof"
  17         "github.com/dustin/go-humanize"
  18         "github.com/jessevdk/go-flags"
  19
  20         "github.com/anacrolix/torrent"
  21         "github.com/anacrolix/torrent/metainfo"
  22 )
  23
  24 // fmt.Fprintf(os.Stderr, "Usage: %s \n", os.Args[0])
  25
  26 func resolvedPeerAddrs(ss []string) (ret []torrent.PeerInfo, err error) {
  27         for _, s := range ss {
  28                 var addr *net.TCPAddr
  29                 addr, err = net.ResolveTCPAddr("tcp", s)
  30                 if err != nil {
  31                         return
  32                 }
  33                 ret = append(ret, torrent.PeerInfo{
  34                         Addr: addr,
  35                 })
  36         }
  37         return
  38 }
  39
  40 func bytesCompleted(tc *torrent.Client) (ret int64) {
  41         for _, t := range tc.Torrents() {
  42                 if t.Info() != nil {
  43                         ret += t.BytesCompleted()
  44                 }
  45         }
  46         return
  47 }
  48
  49 // Returns an estimate of the total bytes for all torrents.
  50 func totalBytesEstimate(tc *torrent.Client) (ret int64) {
  51         var noInfo, hadInfo int64
  52         for _, t := range tc.Torrents() {
  53                 info := t.Info()
  54                 if info == nil {
  55                         noInfo++
  56                         continue
  57                 }
  58                 ret += info.TotalLength()
  59                 hadInfo++
  60         }
  61         if hadInfo != 0 {
  62                 // Treat each torrent without info as the average of those with,
  63                 // rounded up.
  64                 ret += (noInfo*ret + hadInfo - 1) / hadInfo
  65         }
  66         return
  67 }
  68
  69 func progressLine(tc *torrent.Client) string {
  70         return fmt.Sprintf("\033[K%s / %s\r", humanize.Bytes(uint64(bytesCompleted(tc))), humanize.Bytes(uint64(totalBytesEstimate(tc))))
  71 }
  72
  73 func dstFileName(picked string) string {
  74         parts := strings.Split(picked, "/")
  75         return parts[len(parts)-1]
  76 }
  77
  78 func main() {
  79         log.SetFlags(log.LstdFlags | log.Lshortfile)
  80         rootGroup := struct {
  81                 Client    *torrent.ClientConfig `group:"Client Options"`
  82                 TestPeers []string              `long:"test-peer" description:"address of peer to inject to every torrent"`
  83                 Pick      string                `long:"pick" description:"filename to pick"`
  84         }{
  85                 Client: torrent.NewDefaultClientConfig(),
  86         }
  87         // Don't pass flags.PrintError because it's inconsistent with printing.
  88         // https://github.com/jessevdk/go-flags/issues/132
  89         parser := flags.NewParser(&rootGroup, flags.HelpFlag|flags.PassDoubleDash)
  90         parser.Usage = "[OPTIONS] (magnet URI or .torrent file path)..."
  91         posArgs, err := parser.Parse()
  92         if err != nil {
  93                 fmt.Fprintf(os.Stderr, "%s", "Download from the BitTorrent network.\n\n")
  94                 fmt.Println(err)
  95                 os.Exit(2)
  96         }
  97         log.Printf("File to pick: %s", rootGroup.Pick)
  98
  99         testPeers, err := resolvedPeerAddrs(rootGroup.TestPeers)
 100         if err != nil {
 101                 log.Fatal(err)
 102         }
 103
 104         if len(posArgs) == 0 {
 105                 fmt.Fprintln(os.Stderr, "no torrents specified")
 106                 return
 107         }
 108
 109         tmpdir, err := os.MkdirTemp("", "torrent-pick-")
 110         if err != nil {
 111                 log.Fatal(err)
 112         }
 113
 114         defer os.RemoveAll(tmpdir)
 115
 116         rootGroup.Client.DataDir = tmpdir
 117
 118         client, err := torrent.NewClient(rootGroup.Client)
 119         if err != nil {
 120                 log.Fatalf("error creating client: %s", err)
 121         }
 122         http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
 123                 client.WriteStatus(w)
 124         })
 125         defer client.Close()
 126
 127         dstName := dstFileName(rootGroup.Pick)
 128
 129         f, err := os.Create(dstName)
 130         if err != nil {
 131                 log.Fatal(err)
 132         }
 133         dstWriter := bufio.NewWriter(f)
 134
 135         done := make(chan struct{})
 136         for _, arg := range posArgs {
 137                 t := func() *torrent.Torrent {
 138                         if strings.HasPrefix(arg, "magnet:") {
 139                                 t, err := client.AddMagnet(arg)
 140                                 if err != nil {
 141                                         log.Fatalf("error adding magnet: %s", err)
 142                                 }
 143                                 return t
 144                         } else {
 145                                 metaInfo, err := metainfo.LoadFromFile(arg)
 146                                 if err != nil {
 147                                         log.Fatal(err)
 148                                 }
 149                                 t, err := client.AddTorrent(metaInfo)
 150                                 if err != nil {
 151                                         log.Fatal(err)
 152                                 }
 153                                 return t
 154                         }
 155                 }()
 156                 t.AddPeers(testPeers)
 157
 158                 go func() {
 159                         defer close(done)
 160                         <-t.GotInfo()
 161                         for _, file := range t.Files() {
 162                                 if file.DisplayPath() != rootGroup.Pick {
 163                                         continue
 164                                 }
 165                                 file.Download()
 166                                 srcReader := file.NewReader()
 167                                 defer srcReader.Close()
 168                                 io.Copy(dstWriter, srcReader)
 169                                 return
 170                         }
 171                         log.Print("file not found")
 172                 }()
 173         }
 174
 175         ticker := time.NewTicker(time.Second)
 176         defer ticker.Stop()
 177 waitDone:
 178         for {
 179                 select {
 180                 case <-done:
 181                         break waitDone
 182                 case <-ticker.C:
 183                         os.Stdout.WriteString(progressLine(client))
 184                 }
 185         }
 186         if rootGroup.Client.Seed {
 187                 select {}
 188         }
 189 }