Skip to content

Commit

Permalink
update: init multi threading
Browse files Browse the repository at this point in the history
  • Loading branch information
Esonhugh committed Mar 19, 2024
1 parent b098572 commit ec0b8fe
Show file tree
Hide file tree
Showing 11 changed files with 213 additions and 20 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,6 @@ build:

build-static:
GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o $(BUILD_DIR)/$(MAIN_PROGRAM_NAME)-linux-static main.go

upx $(BUILD_DIR)/$(MAIN_PROGRAM_NAME)-linux-static
clean:
rm -rf $(BUILD_DIR)
34 changes: 27 additions & 7 deletions cmd/all/all.go
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
package all

import (
"net"
"os"

command "github.com/esonhugh/k8spider/cmd"
"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg"
"github.com/esonhugh/k8spider/pkg/mutli"
"github.com/esonhugh/k8spider/pkg/scanner"
"github.com/miekg/dns"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
Expand All @@ -29,18 +32,35 @@ var AllCmd = &cobra.Command{
log.Warnf("ParseStringToIPNet failed: %v", err)
return
}
var records define.Records = pkg.ScanSubnet(ipNets)
if records == nil || len(records) == 0 {
log.Warnf("ScanSubnet Found Nothing: %v", err)
return
if command.Opts.BatchMode {
RunBatch(ipNets)
} else {
Run(ipNets)
}
records = pkg.ScanSvcForPorts(records)
printResult(records)
records = pkg.DumpAXFR(dns.Fqdn(command.Opts.Zone), "ns.dns."+command.Opts.Zone+":53")
records := scanner.DumpAXFR(dns.Fqdn(command.Opts.Zone), "ns.dns."+command.Opts.Zone+":53")
printResult(records)
},
}

func Run(net *net.IPNet) {
var records define.Records = scanner.ScanSubnet(net)
if records == nil || len(records) == 0 {
log.Warnf("ScanSubnet Found Nothing")
return
}
records = scanner.ScanSvcForPorts(records)
printResult(records)
}

func RunBatch(net *net.IPNet) {
scan := mutli.ScanAll(net)
var finalRecord []define.Record
for r := range scan {
finalRecord = append(finalRecord, r...)
}
printResult(finalRecord)
}

func printResult(records define.Records) {
if command.Opts.OutputFile != "" {
f, err := os.OpenFile(command.Opts.OutputFile, os.O_CREATE|os.O_WRONLY, 0644)
Expand Down
4 changes: 2 additions & 2 deletions cmd/axfr/axfr.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (

command "github.com/esonhugh/k8spider/cmd"
"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg"
"github.com/esonhugh/k8spider/pkg/scanner"
"github.com/miekg/dns"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
Expand Down Expand Up @@ -35,7 +35,7 @@ var AxfrCmd = &cobra.Command{
}

log.Debugf("same command: dig axfr %v @%v", zone, dnsServer)
var records define.Records = pkg.DumpAXFR(zone, dnsServer)
var records define.Records = scanner.DumpAXFR(zone, dnsServer)
if command.Opts.OutputFile != "" {
f, err := os.OpenFile(command.Opts.OutputFile, os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
Expand Down
9 changes: 6 additions & 3 deletions cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@ var Opts = struct {
SvcDomains []string
Zone string
OutputFile string
LogLevel string
Verbose string

BatchMode bool
}{}

func init() {
Expand All @@ -26,15 +28,16 @@ func init() {
RootCmd.PersistentFlags().StringSliceVarP(&Opts.SvcDomains, "svc-domains", "s", []string{}, "service domains, like: kubernetes.default,etcd.default don't add zone like svc.cluster.local")
RootCmd.PersistentFlags().StringVarP(&Opts.Zone, "zone", "z", "cluster.local", "zone")
RootCmd.PersistentFlags().StringVarP(&Opts.OutputFile, "output-file", "o", "", "output file")
RootCmd.PersistentFlags().StringVarP(&Opts.LogLevel, "log-level", "l", "info", "log level")
RootCmd.PersistentFlags().StringVarP(&Opts.Verbose, "verbose", "v", "info", "log level (debug,info,trace,warn,error,fatal,panic)")
RootCmd.PersistentFlags().BoolVarP(&Opts.BatchMode, "batch-mode", "b", false, "batch mode")
}

var RootCmd = &cobra.Command{
Use: "k8spider",
Short: "k8spider is a tool to discover k8s services",
Long: "k8spider is a tool to discover k8s services",
PersistentPreRun: func(cmd *cobra.Command, args []string) {
SetLogLevel(Opts.LogLevel)
SetLogLevel(Opts.Verbose)
if Opts.DnsServer != "" {
pkg.NetResolver = &net.Resolver{
PreferGo: true,
Expand Down
4 changes: 2 additions & 2 deletions cmd/service/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (

command "github.com/esonhugh/k8spider/cmd"
"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg"
"github.com/esonhugh/k8spider/pkg/scanner"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
)
Expand All @@ -27,7 +27,7 @@ var ServiceCmd = &cobra.Command{
for _, domain := range command.Opts.SvcDomains {
records = append(records, define.Record{SvcDomain: fmt.Sprintf("%s.svc.%s", domain, command.Opts.Zone)})
}
records = pkg.ScanSvcForPorts(records)
records = scanner.ScanSvcForPorts(records)
if command.Opts.OutputFile != "" {
f, err := os.OpenFile(command.Opts.OutputFile, os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
Expand Down
3 changes: 2 additions & 1 deletion cmd/subnet/subnet.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
command "github.com/esonhugh/k8spider/cmd"
"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg"
"github.com/esonhugh/k8spider/pkg/scanner"
log "github.com/sirupsen/logrus"
"github.com/spf13/cobra"
)
Expand All @@ -27,7 +28,7 @@ var SubNetCmd = &cobra.Command{
log.Warnf("ParseStringToIPNet failed: %v", err)
return
}
var records define.Records = pkg.ScanSubnet(ipNets)
var records define.Records = scanner.ScanSubnet(ipNets)
if records == nil || len(records) == 0 {
log.Warnf("ScanSubnet Found Nothing: %v", err)
return
Expand Down
13 changes: 13 additions & 0 deletions pkg/mutli/executor.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
package mutli

import (
"net"

"github.com/esonhugh/k8spider/define"
)

func ScanAll(subnet *net.IPNet) (result <-chan []define.Record) {
subs := NewSubnetScanner()
result = ScanServiceWithChan(subs.ScanSubnet(subnet))
return result
}
47 changes: 47 additions & 0 deletions pkg/mutli/subnet.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
package mutli

import (
"net"
"sync"

"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg"
"github.com/esonhugh/k8spider/pkg/scanner"
log "github.com/sirupsen/logrus"
)

type SubnetScanner struct {
wg *sync.WaitGroup
}

func NewSubnetScanner() *SubnetScanner {
return &SubnetScanner{
wg: new(sync.WaitGroup),
}
}

func (s *SubnetScanner) ScanSubnet(subnet *net.IPNet) <-chan []define.Record {
if subnet == nil {
log.Tracef("subnet is nil")
return nil
}
out := make(chan []define.Record, 100)
go func() {
if subnets, err := pkg.SubnetShift(subnet, 4); err != nil {
go s.scan(subnet, out)
} else {
for _, sn := range subnets {
go s.scan(sn, out)
}
}
s.wg.Wait()
close(out)
}()
return out
}

func (s *SubnetScanner) scan(subnet *net.IPNet, to chan []define.Record) {
s.wg.Add(1)
to <- scanner.ScanSubnet(subnet)
s.wg.Done()
}
17 changes: 17 additions & 0 deletions pkg/mutli/svc.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
package mutli

import (
"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg/scanner"
)

func ScanServiceWithChan(rev <-chan []define.Record) <-chan []define.Record {
out := make(chan []define.Record, 100)
go func() {
for records := range rev {
out <- scanner.ScanSvcForPorts(records)
}
close(out)
}()
return out
}
9 changes: 5 additions & 4 deletions pkg/scanner.go → pkg/scanner/scanner.go
Original file line number Diff line number Diff line change
@@ -1,17 +1,18 @@
package pkg
package scanner

import (
"net"
"strings"

"github.com/esonhugh/k8spider/define"
"github.com/esonhugh/k8spider/pkg"
"github.com/miekg/dns"
log "github.com/sirupsen/logrus"
)

func ScanSubnet(subnet *net.IPNet) (records []define.Record) {
for _, ip := range ParseIPNetToIPs(subnet) {
ptr := PTRRecord(ip)
for _, ip := range pkg.ParseIPNetToIPs(subnet) {
ptr := pkg.PTRRecord(ip)
if len(ptr) > 0 {
for _, domain := range ptr {
log.Infof("PTRrecord %v --> %v", ip, domain)
Expand All @@ -27,7 +28,7 @@ func ScanSubnet(subnet *net.IPNet) (records []define.Record) {

func ScanSvcForPorts(records []define.Record) []define.Record {
for i, r := range records {
cname, srv, err := SRVRecord(r.SvcDomain)
cname, srv, err := pkg.SRVRecord(r.SvcDomain)
if err != nil {
log.Debugf("SRVRecord for %v,failed: %v", r.SvcDomain, err)
continue
Expand Down
91 changes: 91 additions & 0 deletions pkg/subnets.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
package pkg

import (
"fmt"
"math"
"net"
)

// SubnetInto wraps SubnetShift and divides a network into at least count-many,
// equal-sized subnets, which are as large as allowed.
func SubnetInto(network *net.IPNet, count int) ([]*net.IPNet, error) {
maskBits, _ := network.Mask.Size()
hostBits := 32 - maskBits
hostCount := 1 << uint(hostBits)

// divide hosts among subnets
ideal := float64(hostCount) / float64(count)
// largest power of 2, not exceeding the ideal (float64 to int conversion
// truncates toward zero)
newHostBits := int(math.Log2(ideal))
shift := hostBits - newHostBits
return SubnetShift(network, shift)
}

// SubnetShift divides a network into subnets by shifting the given number of bits.
func SubnetShift(network *net.IPNet, bits int) ([]*net.IPNet, error) {
if bits < 0 {
return nil, fmt.Errorf("bit shift may not be negative, got %d", bits)
}
if bits > 31 {
return nil, fmt.Errorf("network subnets cannot be divided %d times", bits)
}
// network divides into 2^bits subnets
subnetCount := 1 << uint(bits)
subnets := make([]*net.IPNet, subnetCount)

// network info
start := network.IP
maskBits, _ := network.Mask.Size()
hostBits := 32 - maskBits

if maskBits+bits > 32 {
return nil, fmt.Errorf("network subnet mask greater than /32, /%d is invalid", maskBits+bits)
}

// divide network into subnets
newMaskBits := maskBits + bits
newHostBits := hostBits - bits
// subnet bitmasks are shifted by 'bits' places
newMask := net.CIDRMask(newMaskBits, 32)

// hosts per subnet
hostCount := 1 << uint(newHostBits)

for i := 0; i < subnetCount; i++ {
ip := numeric(start) + uint32(i*hostCount)
subnets[i] = &net.IPNet{
IP: bytewise(ip),
Mask: newMask,
}
}

return subnets, nil
}

// IP <-> integer transforms

// numeric returns a uint32 numeric representation of a net.IP.
func numeric(bytes net.IP) uint32 {
var ip uint32
// most significant to least significant
for i, b := range []byte(bytes) {
// bitwise or ("append" in this case)
ip |= uint32(b) << (8 * uint32(3-i))
}
return ip
}

// bytewise returns a net.IP byte slice alias representation of an uint32.
// Note that not all uint32 values are valid IP addresses.
func bytewise(numeric uint32) net.IP {
ip := make([]byte, 4)
// least significant to most significant
for i := 3; i >= 0; i-- {
// AND away all but least significant
ip[i] = byte(numeric & 0xFF)
// nuke least significant byte
numeric >>= 8
}
return net.IP(ip)
}

0 comments on commit ec0b8fe

Please sign in to comment.