cmd/torrent/main.go

   1 // Downloads torrents from the command-line.
   2 package main
   3
   4 import (
   5         "fmt"
   6         "log"
   7         "net"
   8         "net/http"
   9         _ "net/http/pprof"
  10         "os"
  11         "strings"
  12         "time"
  13
  14         _ "github.com/anacrolix/envpprof"
  15         "github.com/anacrolix/tagflag"
  16         "github.com/dustin/go-humanize"
  17         "github.com/gosuri/uiprogress"
  18
  19         "github.com/anacrolix/torrent"
  20         "github.com/anacrolix/torrent/metainfo"
  21         "github.com/anacrolix/torrent/storage"
  22 )
  23
  24 func resolvedPeerAddrs(ss []string) (ret []torrent.Peer, err error) {
  25         for _, s := range ss {
  26                 var addr *net.TCPAddr
  27                 addr, err = net.ResolveTCPAddr("tcp", s)
  28                 if err != nil {
  29                         return
  30                 }
  31                 ret = append(ret, torrent.Peer{
  32                         IP:   addr.IP,
  33                         Port: addr.Port,
  34                 })
  35         }
  36         return
  37 }
  38
  39 func torrentBar(t torrent.Torrent) {
  40         bar := uiprogress.AddBar(1)
  41         bar.AppendCompleted()
  42         bar.AppendFunc(func(*uiprogress.Bar) (ret string) {
  43                 select {
  44                 case <-t.GotInfo():
  45                 default:
  46                         return "getting info"
  47                 }
  48                 if t.Seeding() {
  49                         return "seeding"
  50                 } else if t.BytesCompleted() == t.Info().TotalLength() {
  51                         return "completed"
  52                 } else {
  53                         return fmt.Sprintf("downloading (%s/%s)", humanize.Bytes(uint64(t.BytesCompleted())), humanize.Bytes(uint64(t.Info().TotalLength())))
  54                 }
  55         })
  56         bar.PrependFunc(func(*uiprogress.Bar) string {
  57                 return t.Name()
  58         })
  59         go func() {
  60                 <-t.GotInfo()
  61                 bar.Total = int(t.Info().TotalLength())
  62                 for {
  63                         bc := t.BytesCompleted()
  64                         bar.Set(int(bc))
  65                         time.Sleep(time.Second)
  66                 }
  67         }()
  68 }
  69
  70 func addTorrents(client *torrent.Client) {
  71         for _, arg := range opts.Torrent {
  72                 t := func() torrent.Torrent {
  73                         if strings.HasPrefix(arg, "magnet:") {
  74                                 t, err := client.AddMagnet(arg)
  75                                 if err != nil {
  76                                         log.Fatalf("error adding magnet: %s", err)
  77                                 }
  78                                 return t
  79                         } else {
  80                                 metaInfo, err := metainfo.LoadFromFile(arg)
  81                                 if err != nil {
  82                                         fmt.Fprintf(os.Stderr, "error loading torrent file %q: %s\n", arg, err)
  83                                         os.Exit(1)
  84                                 }
  85                                 t, err := client.AddTorrent(metaInfo)
  86                                 if err != nil {
  87                                         log.Fatal(err)
  88                                 }
  89                                 return t
  90                         }
  91                 }()
  92                 torrentBar(t)
  93                 err := t.AddPeers(func() (ret []torrent.Peer) {
  94                         for _, ta := range opts.TestPeer {
  95                                 ret = append(ret, torrent.Peer{
  96                                         IP:   ta.IP,
  97                                         Port: ta.Port,
  98                                 })
  99                         }
 100                         return
 101                 }())
 102                 if err != nil {
 103                         log.Fatal(err)
 104                 }
 105                 go func() {
 106                         <-t.GotInfo()
 107                         t.DownloadAll()
 108                 }()
 109         }
 110 }
 111
 112 var opts struct {
 113         torrent.Config `name:"Client"`
 114         Mmap           bool           `help:"memory-map torrent data"`
 115         TestPeer       []*net.TCPAddr `short:"p" help:"addresses of some starting peers"`
 116         Torrent        []string       `type:"pos" arity:"+" help:"torrent file path or magnet uri"`
 117 }
 118
 119 func main() {
 120         log.SetFlags(log.LstdFlags | log.Lshortfile)
 121         tagflag.Parse(&opts, tagflag.SkipBadTypes())
 122         clientConfig := opts.Config
 123         if opts.Mmap {
 124                 clientConfig.DefaultStorage = storage.NewMMap("")
 125         }
 126
 127         client, err := torrent.NewClient(&clientConfig)
 128         if err != nil {
 129                 log.Fatalf("error creating client: %s", err)
 130         }
 131         defer client.Close()
 132         http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
 133                 client.WriteStatus(w)
 134         })
 135         uiprogress.Start()
 136         addTorrents(client)
 137         if client.WaitAll() {
 138                 log.Print("downloaded ALL the torrents")
 139         } else {
 140                 log.Fatal("y u no complete torrents?!")
 141         }
 142         if opts.Seed {
 143                 select {}
 144         }
 145 }