dht_hot_files.go

dht_hot_files.go

package main

import (
	"fmt"
	"io/ioutil"
	"net"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/shiyanhui/dht"
	log "github.com/sirupsen/logrus"
)

type Message struct {
	InfoHash string `json:"info_hash"`
	IP       net.IP `json:"ip"`
	Time     int64  `json:"time"`
}

func rankMapStringInt(values map[string]int) []string {
	type kv struct {
		Key   string
		Value int
	}
	var ss []kv
	for k, v := range values {
		ss = append(ss, kv{k, v})
	}
	sort.Slice(ss, func(i, j int) bool {
		return ss[i].Value > ss[j].Value
	})
	ranked := make([]string, len(values))
	for i, kv := range ss {
		ranked[i] = kv.Key
	}
	return ranked
}

func main() {
	config := dht.NewStandardConfig()
	config.Address = ":5643"
	config.MaxNodes = 655350
	d := dht.New(config)
	log.SetLevel(log.DebugLevel)
	counter := make(map[string]int)
	d.OnGetPeersResponse = func(infoHash string, peer *dht.Peer) {
		m := Message{
			InfoHash: infoHash,
			IP:       peer.IP,
			Time:     time.Now().Unix(),
		}
		item, ok := counter[m.InfoHash]
		if !ok {
			item = 1
		} else {
			item = item + 1
		}
		counter[m.InfoHash] = item
		log.Debugf("get %x from %s", infoHash, peer.IP)
	}
	items, _ := ioutil.ReadFile("/tmp/input.txt")
	data := strings.Split(string(items), "\n")
	wg := sync.WaitGroup{}
	for _, infoHash := range data {
		go func(infoHash string, group *sync.WaitGroup) {
			defer group.Done()
			defer func() {
				if r := recover(); r != nil {
					fmt.Println("Recovered in f", r)
				}
			}()
			for {
				err := d.GetPeers(strings.ToLower(infoHash))
				if err != nil && err != dht.ErrNotReady {
					log.Fatal(err)
				}
				if err == dht.ErrNotReady {
					//log.Debug("waiting dht server ready...")
					time.Sleep(time.Second * 1)
					continue
				}
				break
			}
		}(infoHash, &wg)
		wg.Add(1)
	}
	go d.Run()
	wg.Wait()
	log.Debug("waiting 10 seconds")
	time.Sleep(time.Second * 120)
	var max = 10
	var count = 0
	for _, key := range rankMapStringInt(counter) {
		count += 1
		if count > max {
			break
		}
		fmt.Printf("%x,%d\n", key, counter[key])
	}
	log.Debug("done")
}