cmd/dht-get-peers/main.go

   1 package main
   2
   3 import (
   4         "flag"
   5         "fmt"
   6         "io"
   7         "log"
   8         "net"
   9         "os"
  10         "os/signal"
  11
  12         _ "github.com/anacrolix/envpprof"
  13
  14         "github.com/anacrolix/torrent/dht"
  15         _ "github.com/anacrolix/torrent/util/profile"
  16 )
  17
  18 var (
  19         tableFileName = flag.String("tableFile", "", "name of file for storing node info")
  20         serveAddr     = flag.String("serveAddr", ":0", "local UDP address")
  21         infoHash      = flag.String("infoHash", "", "torrent infohash")
  22         once          = flag.Bool("once", false, "only do one scrape iteration")
  23
  24         s        *dht.Server
  25         quitting = make(chan struct{})
  26 )
  27
  28 func loadTable() error {
  29         if *tableFileName == "" {
  30                 return nil
  31         }
  32         f, err := os.Open(*tableFileName)
  33         if os.IsNotExist(err) {
  34                 return nil
  35         }
  36         if err != nil {
  37                 return fmt.Errorf("error opening table file: %s", err)
  38         }
  39         defer f.Close()
  40         added := 0
  41         for {
  42                 b := make([]byte, dht.CompactIPv4NodeInfoLen)
  43                 _, err := io.ReadFull(f, b)
  44                 if err == io.EOF {
  45                         break
  46                 }
  47                 if err != nil {
  48                         return fmt.Errorf("error reading table file: %s", err)
  49                 }
  50                 var ni dht.NodeInfo
  51                 err = ni.UnmarshalCompactIPv4(b)
  52                 if err != nil {
  53                         return fmt.Errorf("error unmarshaling compact node info: %s", err)
  54                 }
  55                 s.AddNode(ni)
  56                 added++
  57         }
  58         log.Printf("loaded %d nodes from table file", added)
  59         return nil
  60 }
  61
  62 func init() {
  63         log.SetFlags(log.LstdFlags | log.Lshortfile)
  64         flag.Parse()
  65         switch len(*infoHash) {
  66         case 20:
  67         case 40:
  68                 _, err := fmt.Sscanf(*infoHash, "%x", infoHash)
  69                 if err != nil {
  70                         log.Fatal(err)
  71                 }
  72         default:
  73                 log.Fatal("require 20 byte infohash")
  74         }
  75         var err error
  76         s, err = dht.NewServer(&dht.ServerConfig{
  77                 Addr: *serveAddr,
  78         })
  79         if err != nil {
  80                 log.Fatal(err)
  81         }
  82         err = loadTable()
  83         if err != nil {
  84                 log.Fatalf("error loading table: %s", err)
  85         }
  86         log.Printf("dht server on %s, ID is %x", s.Addr(), s.ID())
  87         setupSignals()
  88 }
  89
  90 func saveTable() error {
  91         goodNodes := s.Nodes()
  92         if *tableFileName == "" {
  93                 if len(goodNodes) != 0 {
  94                         log.Print("good nodes were discarded because you didn't specify a table file")
  95                 }
  96                 return nil
  97         }
  98         f, err := os.OpenFile(*tableFileName, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
  99         if err != nil {
 100                 return fmt.Errorf("error opening table file: %s", err)
 101         }
 102         defer f.Close()
 103         for _, nodeInfo := range goodNodes {
 104                 var b [dht.CompactIPv4NodeInfoLen]byte
 105                 err := nodeInfo.PutCompact(b[:])
 106                 if err != nil {
 107                         return fmt.Errorf("error compacting node info: %s", err)
 108                 }
 109                 _, err = f.Write(b[:])
 110                 if err != nil {
 111                         return fmt.Errorf("error writing compact node info: %s", err)
 112                 }
 113         }
 114         log.Printf("saved %d nodes to table file", len(goodNodes))
 115         return nil
 116 }
 117
 118 func setupSignals() {
 119         ch := make(chan os.Signal)
 120         signal.Notify(ch, os.Interrupt)
 121         go func() {
 122                 <-ch
 123                 close(quitting)
 124         }()
 125 }
 126
 127 func main() {
 128         seen := make(map[string]struct{})
 129 getPeers:
 130         for {
 131                 ps, err := s.Announce(*infoHash, 0, false)
 132                 if err != nil {
 133                         log.Fatal(err)
 134                 }
 135         values:
 136                 for {
 137                         select {
 138                         case v, ok := <-ps.Peers:
 139                                 if !ok {
 140                                         break values
 141                                 }
 142                                 log.Printf("received %d peers from %x", len(v.Peers), v.NodeInfo.ID)
 143                                 for _, p := range v.Peers {
 144                                         if _, ok := seen[p.String()]; ok {
 145                                                 continue
 146                                         }
 147                                         seen[p.String()] = struct{}{}
 148                                         fmt.Println((&net.UDPAddr{
 149                                                 IP:   p.IP[:],
 150                                                 Port: int(p.Port),
 151                                         }).String())
 152                                 }
 153                         case <-quitting:
 154                                 break getPeers
 155                         }
 156                 }
 157                 if *once {
 158                         break
 159                 }
 160         }
 161         if err := saveTable(); err != nil {
 162                 log.Printf("error saving node table: %s", err)
 163         }
 164 }