cmd/torrent/main.go

   1 // Downloads torrents from the command-line.
   2 package main
   3
   4 import (
   5         "expvar"
   6         "fmt"
   7         "log"
   8         "net"
   9         "net/http"
  10         "os"
  11         "strings"
  12         "time"
  13
  14         "github.com/anacrolix/torrent/iplist"
  15
  16         "github.com/anacrolix/envpprof"
  17         "github.com/anacrolix/tagflag"
  18         "github.com/dustin/go-humanize"
  19         "github.com/gosuri/uiprogress"
  20         "golang.org/x/time/rate"
  21
  22         "github.com/anacrolix/torrent"
  23         "github.com/anacrolix/torrent/metainfo"
  24         "github.com/anacrolix/torrent/storage"
  25 )
  26
  27 var progress = uiprogress.New()
  28
  29 func torrentBar(t *torrent.Torrent) {
  30         bar := progress.AddBar(1)
  31         bar.AppendCompleted()
  32         bar.AppendFunc(func(*uiprogress.Bar) (ret string) {
  33                 select {
  34                 case <-t.GotInfo():
  35                 default:
  36                         return "getting info"
  37                 }
  38                 if t.Seeding() {
  39                         return "seeding"
  40                 } else if t.BytesCompleted() == t.Info().TotalLength() {
  41                         return "completed"
  42                 } else {
  43                         return fmt.Sprintf("downloading (%s/%s)", humanize.Bytes(uint64(t.BytesCompleted())), humanize.Bytes(uint64(t.Info().TotalLength())))
  44                 }
  45         })
  46         bar.PrependFunc(func(*uiprogress.Bar) string {
  47                 return t.Name()
  48         })
  49         go func() {
  50                 <-t.GotInfo()
  51                 tl := int(t.Info().TotalLength())
  52                 if tl == 0 {
  53                         bar.Set(1)
  54                         return
  55                 }
  56                 bar.Total = tl
  57                 for {
  58                         bc := t.BytesCompleted()
  59                         bar.Set(int(bc))
  60                         time.Sleep(time.Second)
  61                 }
  62         }()
  63 }
  64
  65 func addTorrents(client *torrent.Client) {
  66         for _, arg := range flags.Torrent {
  67                 t := func() *torrent.Torrent {
  68                         if strings.HasPrefix(arg, "magnet:") {
  69                                 t, err := client.AddMagnet(arg)
  70                                 if err != nil {
  71                                         log.Fatalf("error adding magnet: %s", err)
  72                                 }
  73                                 return t
  74                         } else if strings.HasPrefix(arg, "http://") || strings.HasPrefix(arg, "https://") {
  75                                 response, err := http.Get(arg)
  76                                 if err != nil {
  77                                         log.Fatalf("Error downloading torrent file: %s", err)
  78                                 }
  79
  80                                 metaInfo, err := metainfo.Load(response.Body)
  81                                 defer response.Body.Close()
  82                                 if err != nil {
  83                                         fmt.Fprintf(os.Stderr, "error loading torrent file %q: %s\n", arg, err)
  84                                         os.Exit(1)
  85                                 }
  86                                 t, err := client.AddTorrent(metaInfo)
  87                                 if err != nil {
  88                                         log.Fatal(err)
  89                                 }
  90                                 return t
  91                         } else if strings.HasPrefix(arg, "infohash:") {
  92                                 t, _ := client.AddTorrentInfoHash(metainfo.NewHashFromHex(strings.TrimPrefix(arg, "infohash:")))
  93                                 return t
  94                         } else {
  95                                 metaInfo, err := metainfo.LoadFromFile(arg)
  96                                 if err != nil {
  97                                         fmt.Fprintf(os.Stderr, "error loading torrent file %q: %s\n", arg, err)
  98                                         os.Exit(1)
  99                                 }
 100                                 t, err := client.AddTorrent(metaInfo)
 101                                 if err != nil {
 102                                         log.Fatal(err)
 103                                 }
 104                                 return t
 105                         }
 106                 }()
 107                 torrentBar(t)
 108                 t.AddPeers(func() (ret []torrent.Peer) {
 109                         for _, ta := range flags.TestPeer {
 110                                 ret = append(ret, torrent.Peer{
 111                                         IP:   ta.IP,
 112                                         Port: ta.Port,
 113                                 })
 114                         }
 115                         return
 116                 }())
 117                 go func() {
 118                         <-t.GotInfo()
 119                         t.DownloadAll()
 120                 }()
 121         }
 122 }
 123
 124 var flags = struct {
 125         Mmap            bool           `help:"memory-map torrent data"`
 126         TestPeer        []*net.TCPAddr `help:"addresses of some starting peers"`
 127         Seed            bool           `help:"seed after download is complete"`
 128         Addr            *net.TCPAddr   `help:"network listen addr"`
 129         UploadRate      tagflag.Bytes  `help:"max piece bytes to send per second"`
 130         DownloadRate    tagflag.Bytes  `help:"max bytes per second down from peers"`
 131         Debug           bool
 132         PackedBlocklist string
 133         Stats           *bool
 134         tagflag.StartPos
 135         Torrent []string `arity:"+" help:"torrent file path or magnet uri"`
 136 }{
 137         UploadRate:   -1,
 138         DownloadRate: -1,
 139 }
 140
 141 func stdoutAndStderrAreSameFile() bool {
 142         fi1, _ := os.Stdout.Stat()
 143         fi2, _ := os.Stderr.Stat()
 144         return os.SameFile(fi1, fi2)
 145 }
 146
 147 func statsEnabled() bool {
 148         if flags.Stats == nil {
 149                 return flags.Debug
 150         }
 151         return *flags.Stats
 152 }
 153
 154 func main() {
 155         log.SetFlags(log.LstdFlags | log.Lshortfile)
 156         tagflag.Parse(&flags)
 157         defer envpprof.Stop()
 158         clientConfig := torrent.NewDefaultClientConfig()
 159         clientConfig.Debug = flags.Debug
 160         clientConfig.Seed = flags.Seed
 161         if flags.PackedBlocklist != "" {
 162                 blocklist, err := iplist.MMapPackedFile(flags.PackedBlocklist)
 163                 if err != nil {
 164                         log.Fatalf("error loading blocklist: %s", err)
 165                 }
 166                 defer blocklist.Close()
 167                 clientConfig.IPBlocklist = blocklist
 168         }
 169         if flags.Mmap {
 170                 clientConfig.DefaultStorage = storage.NewMMap("")
 171         }
 172         if flags.Addr != nil {
 173                 clientConfig.SetListenAddr(flags.Addr.String())
 174         }
 175         if flags.UploadRate != -1 {
 176                 clientConfig.UploadRateLimiter = rate.NewLimiter(rate.Limit(flags.UploadRate), 256<<10)
 177         }
 178         if flags.DownloadRate != -1 {
 179                 clientConfig.DownloadRateLimiter = rate.NewLimiter(rate.Limit(flags.DownloadRate), 1<<20)
 180         }
 181
 182         client, err := torrent.NewClient(clientConfig)
 183         if err != nil {
 184                 log.Fatalf("error creating client: %s", err)
 185         }
 186         defer client.Close()
 187         // Write status on the root path on the default HTTP muxer. This will be
 188         // bound to localhost somewhere if GOPPROF is set, thanks to the envpprof
 189         // import.
 190         http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
 191                 client.WriteStatus(w)
 192         })
 193         if stdoutAndStderrAreSameFile() {
 194                 log.SetOutput(progress.Bypass())
 195         }
 196         progress.Start()
 197         addTorrents(client)
 198         if client.WaitAll() {
 199                 log.Print("downloaded ALL the torrents")
 200         } else {
 201                 log.Fatal("y u no complete torrents?!")
 202         }
 203         if flags.Seed {
 204                 outputStats(client)
 205                 select {}
 206         }
 207         outputStats(client)
 208 }
 209
 210 func outputStats(cl *torrent.Client) {
 211         if !statsEnabled() {
 212                 return
 213         }
 214         expvar.Do(func(kv expvar.KeyValue) {
 215                 fmt.Printf("%s: %s\n", kv.Key, kv.Value)
 216         })
 217         cl.WriteStatus(os.Stdout)
 218 }