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