probe/discovery/protocol/snmp/snmp.go

package snmp

import (
	"encoding/json"
	"fmt"
	"log"
	"net"
	"strconv"
	"sync"
	"time"

	omcc "git.loafle.net/overflow/model/config/credential"
	omd "git.loafle.net/overflow/model/discovery"
	omm "git.loafle.net/overflow/model/meta"
	omu "git.loafle.net/overflow/model/util"
	"git.loafle.net/overflow_scanner/probe/discovery/session"
	"github.com/k-sone/snmpgo"
)

const (
	defaultPort      = 161
	defaultTimeout   = 1
	defaultCommunity = "public"
)

var (
	defaultOIDs = []string{
		"1.3.6.1.2.1.1.5.0", //sysName
	}

	defaultCredentials = []*omcc.SNMPCredential{
		&omcc.SNMPCredential{
			Version:   "2c",
			Community: "public",
			Port:      json.Number(strconv.Itoa(defaultPort)),
			Timeout:   json.Number(strconv.Itoa(defaultTimeout)),
		},
	}
)

func Scan(discoverySession session.DiscoverySession) error {
	targetHosts := discoverySession.TargetHosts()
	if nil == targetHosts || 0 == len(targetHosts) {
		return nil
	}

	var cs []*omcc.SNMPCredential

	if _csMap := discoverySession.DiscoveryConfig().Credentials; nil != _csMap {
		_cs, ok := _csMap["SNMP"]
		if ok {
			cs = _cs.([]*omcc.SNMPCredential)
		}
	}
	if nil == cs {
		cs = defaultCredentials
	}

	credentials := make(map[string][]*omcc.SNMPCredential)
	for _, _c := range cs {
		credentials[_c.Version] = append(credentials[_c.Version], _c)
	}

	var wg sync.WaitGroup

	_2cCS, ok := credentials["2c"]
	if ok {
		wg.Add(1)
		go func() {
			defer wg.Done()

			for _, target := range targetHosts {
				wg.Add(1)
				go func(target net.IP) {
					defer wg.Done()
					for _, c := range _2cCS {
						if scanV2(target, discoverySession, c) {
							return
						}
					}
				}(target)
			}
		}()
	}

	wg.Wait()

	return nil
}

func scanV2(target net.IP, discoverySession session.DiscoverySession, credential *omcc.SNMPCredential) bool {

	address := fmt.Sprintf("%s:%s", target.String(), credential.Port.String())
	timeout, _ := credential.Timeout.Int64()
	snmp, err := snmpgo.NewSNMP(snmpgo.SNMPArguments{
		Version:   snmpgo.V2c,
		Address:   address,
		Timeout:   time.Second * time.Duration(timeout),
		Retries:   1,
		Community: credential.Community,
	})

	if err != nil {
		// ch <- &SNMPResponse{host, nil, err}
		return false
	}
	defer snmp.Close()

	oids, err := snmpgo.NewOids(defaultOIDs)
	if err != nil {
		// ch <- &SNMPResponse{host, nil, err}
		return false
	}

	pdu, err := snmp.GetRequest(oids)
	if err != nil {
		// ch <- &SNMPResponse{host, nil, err}
		return false
	}

	if pdu.ErrorStatus() != snmpgo.NoError {
		// ch <- &SNMPResponse{host, nil, fmt.Errorf("%s", pdu.ErrorStatus().String())}
		return false
	}

	if pdu == nil {
		// ch <- &SNMPResponse{host, nil, fmt.Errorf("%s", "Empty PDU")}
		return false
	}

	meta := make(map[string]string)
	for _, val := range pdu.VarBinds() {
		meta[val.Oid.String()] = val.Variable.String()
	}

	h := discoverySession.AddHost(&omd.Host{
		MetaIPType:     discoverySession.Zone().MetaIPType,
		Name:           "",
		Address:        target.String(),
		Meta:           meta,
		Zone:           discoverySession.Zone(),
		DiscoveredDate: omu.NowPtr(),
	})

	p := discoverySession.AddPort(&omd.Port{
		MetaPortType: omm.ToMetaPortType(omm.MetaPortTypeEnumUDP),
		PortNumber:   credential.Port,
		Meta:         meta,
		Host:         h,
	})

	s := discoverySession.AddService(&omd.Service{
		MetaCryptoType: omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumNONE),
		Key:            "SNMP",
		Name:           "SNMP V2c",
		Port:           p,
	})

	log.Printf("Host: %v, Port: %v, Service: %v", h, p, s)

	return true
}