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         "io/ioutil"
   9         "log"
  10         "net"
  11         "net/http"
  12         _ "net/http/pprof"
  13         "os"
  14         "strings"
  15         "time"
  16
  17         _ "github.com/anacrolix/envpprof"
  18         "github.com/dustin/go-humanize"
  19         "github.com/jessevdk/go-flags"
  20
  21         "github.com/anacrolix/torrent"
  22         "github.com/anacrolix/torrent/metainfo"
  23 )
  24
  25 // fmt.Fprintf(os.Stderr, "Usage: %s \n", os.Args[0])
  26
  27 func resolvedPeerAddrs(ss []string) (ret []torrent.Peer, err error) {
  28         for _, s := range ss {
  29                 var addr *net.TCPAddr
  30                 addr, err = net.ResolveTCPAddr("tcp", s)
  31                 if err != nil {
  32                         return
  33                 }
  34                 ret = append(ret, torrent.Peer{
  35                         IP:   addr.IP,
  36                         Port: addr.Port,
  37                 })
  38         }
  39         return
  40 }
  41
  42 func bytesCompleted(tc *torrent.Client) (ret int64) {
  43         for _, t := range tc.Torrents() {
  44                 if t.Info() != nil {
  45                         ret += t.BytesCompleted()
  46                 }
  47         }
  48         return
  49 }
  50
  51 // Returns an estimate of the total bytes for all torrents.
  52 func totalBytesEstimate(tc *torrent.Client) (ret int64) {
  53         var noInfo, hadInfo int64
  54         for _, t := range tc.Torrents() {
  55                 info := t.Info()
  56                 if info == nil {
  57                         noInfo++
  58                         continue
  59                 }
  60                 ret += info.TotalLength()
  61                 hadInfo++
  62         }
  63         if hadInfo != 0 {
  64                 // Treat each torrent without info as the average of those with,
  65                 // rounded up.
  66                 ret += (noInfo*ret + hadInfo - 1) / hadInfo
  67         }
  68         return
  69 }
  70
  71 func progressLine(tc *torrent.Client) string {
  72         return fmt.Sprintf("\033[K%s / %s\r", humanize.Bytes(uint64(bytesCompleted(tc))), humanize.Bytes(uint64(totalBytesEstimate(tc))))
  73 }
  74
  75 func dstFileName(picked string) string {
  76         parts := strings.Split(picked, "/")
  77         return parts[len(parts)-1]
  78 }
  79
  80 func main() {
  81         log.SetFlags(log.LstdFlags | log.Lshortfile)
  82         var rootGroup struct {
  83                 Client    torrent.Config `group:"Client Options"`
  84                 TestPeers []string       `long:"test-peer" description:"address of peer to inject to every torrent"`
  85                 Pick      string         `long:"pick" description:"filename to pick"`
  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 := ioutil.TempDir("", "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 waitDone:
 177         for {
 178                 select {
 179                 case <-done:
 180                         break waitDone
 181                 case <-ticker.C:
 182                         os.Stdout.WriteString(progressLine(client))
 183                 }
 184         }
 185         if rootGroup.Client.Seed {
 186                 select {}
 187         }
 188 }