cmd/devp2p: WIP nodeset filter

cmd/devp2p/nodeset.go

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"os"
 	"sort"
 	"time"
 
@@ -55,6 +56,10 @@ func writeNodesJSON(file string, nodes nodeSet) {
 	if err != nil {
 		exit(err)
 	}
+	if file == "-" {
+		os.Stdout.Write(nodesJSON)
+		return
+	}
 	if err := ioutil.WriteFile(file, nodesJSON, 0644); err != nil {
 		exit(err)
 	}

cmd/devp2p/nodesetcmd.go

@@ -18,7 +18,12 @@ package main
 
 import (
 	"fmt"
+	"net"
+	"time"
 
+	"github.com/ethereum/go-ethereum/core/forkid"
+	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/ethereum/go-ethereum/rlp"
 	"gopkg.in/urfave/cli.v1"
 )
 
@@ -28,6 +33,7 @@ var (
 		Usage: "Node set tools",
 		Subcommands: []cli.Command{
 			nodesetInfoCommand,
+			nodesetFilterCommand,
 		},
 	}
 	nodesetInfoCommand = cli.Command{
@@ -36,6 +42,14 @@ var (
 		Action:    nodesetInfo,
 		ArgsUsage: "<nodes.json>",
 	}
+	nodesetFilterCommand = cli.Command{
+		Name:      "filter",
+		Usage:     "Filters a node set",
+		Action:    nodesetFilter,
+		ArgsUsage: "<nodes.json> filters..",
+
+		SkipFlagParsing: true,
+	}
 )
 
 func nodesetInfo(ctx *cli.Context) error {
@@ -47,3 +61,130 @@ func nodesetInfo(ctx *cli.Context) error {
 	fmt.Printf("Set contains %d nodes.\n", len(ns))
 	return nil
 }
+
+func nodesetFilter(ctx *cli.Context) error {
+	if ctx.NArg() < 1 {
+		return fmt.Errorf("need nodes file as argument")
+	}
+	ns := loadNodesJSON(ctx.Args().First())
+	filter, err := andFilter(ctx.Args().Tail())
+	if err != nil {
+		return err
+	}
+
+	result := make(nodeSet)
+	for id, n := range ns {
+		if filter(n) {
+			result[id] = n
+		}
+	}
+	writeNodesJSON("-", result)
+	return nil
+}
+
+type nodeFilter func(nodeJSON) bool
+
+type nodeFilterC struct {
+	narg int
+	fn   func([]string) (nodeFilter, error)
+}
+
+var filterFlags = map[string]nodeFilterC{
+	"-ip":          {1, ipFilter},
+	"-min-age":     {1, minAgeFilter},
+	"-eth-network": {1, ethFilter},
+	"-les-server":  {0, lesFilter},
+}
+
+func parseFilters(args []string) ([]nodeFilter, error) {
+	var filters []nodeFilter
+	for len(args) > 0 {
+		fc, ok := filterFlags[args[0]]
+		if !ok {
+			return nil, fmt.Errorf("invalid filter %q", args[0])
+		}
+		if len(args) < fc.narg {
+			return nil, fmt.Errorf("filter %q wants %d arguments, have %d", args[0], fc.narg, len(args))
+		}
+		filter, err := fc.fn(args[1:])
+		if err != nil {
+			return nil, fmt.Errorf("%s: %v", args[0], err)
+		}
+		filters = append(filters, filter)
+		args = args[fc.narg+1:]
+	}
+	return filters, nil
+}
+
+func andFilter(args []string) (nodeFilter, error) {
+	checks, err := parseFilters(args)
+	if err != nil {
+		return nil, err
+	}
+	f := func(n nodeJSON) bool {
+		for _, filter := range checks {
+			if !filter(n) {
+				return false
+			}
+		}
+		return true
+	}
+	return f, nil
+}
+
+func ipFilter(args []string) (nodeFilter, error) {
+	_, cidr, err := net.ParseCIDR(args[0])
+	if err != nil {
+		return nil, err
+	}
+	f := func(n nodeJSON) bool { return cidr.Contains(n.N.IP()) }
+	return f, nil
+}
+
+func minAgeFilter(args []string) (nodeFilter, error) {
+	minage, err := time.ParseDuration(args[0])
+	if err != nil {
+		return nil, err
+	}
+	f := func(n nodeJSON) bool {
+		age := n.LastSeen.Sub(n.FirstSeen)
+		return age >= minage
+	}
+	return f, nil
+}
+
+func ethFilter(args []string) (nodeFilter, error) {
+	// var config *params.ChainConfig
+	// switch args[0] {
+	// case "mainnet":
+	// 	config = params.MainnetChainConfig
+	// case "rinkeby":
+	// 	config = params.RinkebyChainConfig
+	// case "goerli":
+	// 	config = params.GoerliChainConfig
+	// default:
+	// 	return nil, fmt.Errorf("unknown network %q", args[0])
+	// }
+
+	forkFilter := forkid.NewFilter(nil)
+	f := func(n nodeJSON) bool {
+		var eth struct {
+			ForkID forkid.ID
+			_      []rlp.RawValue `rlp:"tail"`
+		}
+		if n.N.Load(enr.WithEntry("eth", &eth)) != nil {
+			return false
+		}
+		return forkFilter(eth.ForkID) == nil
+	}
+	return f, nil
+}
+
+func lesFilter(args []string) (nodeFilter, error) {
+	f := func(n nodeJSON) bool {
+		var les struct {
+			_ []rlp.RawValue `rlp:"tail"`
+		}
+		return n.N.Load(enr.WithEntry("les", &les)) == nil
+	}
+}