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/anacrolix/missinggo"
  19         "github.com/dustin/go-humanize"
  20         "github.com/jessevdk/go-flags"
  21
  22         "github.com/anacrolix/torrent"
  23         "github.com/anacrolix/torrent/metainfo"
  24 )
  25
  26 // fmt.Fprintf(os.Stderr, "Usage: %s \n", os.Args[0])
  27
  28 func resolvedPeerAddrs(ss []string) (ret []torrent.Peer, err error) {
  29         for _, s := range ss {
  30                 var addr *net.TCPAddr
  31                 addr, err = net.ResolveTCPAddr("tcp", s)
  32                 if err != nil {
  33                         return
  34                 }
  35                 ret = append(ret, torrent.Peer{
  36                         IP:   addr.IP,
  37                         Port: addr.Port,
  38                 })
  39         }
  40         return
  41 }
  42
  43 func bytesCompleted(tc *torrent.Client) (ret int64) {
  44         for _, t := range tc.Torrents() {
  45                 if t.Info() != nil {
  46                         ret += t.BytesCompleted()
  47                 }
  48         }
  49         return
  50 }
  51
  52 // Returns an estimate of the total bytes for all torrents.
  53 func totalBytesEstimate(tc *torrent.Client) (ret int64) {
  54         var noInfo, hadInfo int64
  55         for _, t := range tc.Torrents() {
  56                 info := t.Info()
  57                 if info == nil {
  58                         noInfo++
  59                         continue
  60                 }
  61                 ret += info.TotalLength()
  62                 hadInfo++
  63         }
  64         if hadInfo != 0 {
  65                 // Treat each torrent without info as the average of those with,
  66                 // rounded up.
  67                 ret += (noInfo*ret + hadInfo - 1) / hadInfo
  68         }
  69         return
  70 }
  71
  72 func progressLine(tc *torrent.Client) string {
  73         return fmt.Sprintf("\033[K%s / %s\r", humanize.Bytes(uint64(bytesCompleted(tc))), humanize.Bytes(uint64(totalBytesEstimate(tc))))
  74 }
  75
  76 func dstFileName(picked string) string {
  77         parts := strings.Split(picked, "/")
  78         return parts[len(parts)-1]
  79 }
  80
  81 func main() {
  82         log.SetFlags(log.LstdFlags | log.Lshortfile)
  83         var rootGroup struct {
  84                 Client    torrent.Config `group:"Client Options"`
  85                 TestPeers []string       `long:"test-peer" description:"address of peer to inject to every torrent"`
  86                 Pick      string         `long:"pick" description:"filename to pick"`
  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                 err := t.AddPeers(testPeers)
 158                 if err != nil {
 159                         log.Fatal(err)
 160                 }
 161
 162                 go func() {
 163                         defer close(done)
 164                         <-t.GotInfo()
 165                         for _, file := range t.Files() {
 166                                 if file.DisplayPath() != rootGroup.Pick {
 167                                         continue
 168                                 }
 169                                 srcReader := missinggo.NewSectionReadSeeker(t.NewReader(), file.Offset(), file.Length())
 170                                 io.Copy(dstWriter, srcReader)
 171                                 return
 172                         }
 173                         log.Print("file not found")
 174                 }()
 175         }
 176
 177         ticker := time.NewTicker(time.Second)
 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 }