hswaw/leasifier/main.go

package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/golang/glog"

	mirko "code.hackerspace.pl/hscloud/go/mirko"
	hpb "code.hackerspace.pl/hscloud/hswaw/proto"
)

type lease struct {
	hardware net.HardwareAddr
	ip       net.IP
	from     time.Time
	to       time.Time
}

const timeFmt = "2006/01/02 15:04:05"

func parseLeases(f io.Reader, now time.Time) ([]lease, error) {
	scanner := bufio.NewScanner(f)

	leases := make(map[string]*lease)

	var curAddr *net.IP
	var curStart *time.Time
	var curEnd *time.Time
	var curHW *net.HardwareAddr

	for scanner.Scan() {
		l := scanner.Text()
		l = strings.TrimSpace(l)
		if len(l) < 1 {
			continue
		}

		if strings.HasPrefix(l, "#") {
			continue
		}

		parts := strings.Fields(l)
		if len(parts) < 1 {
			continue
		}

		if curAddr == nil {
			switch parts[0] {
			case "lease":
				if len(parts) < 3 || parts[2] != "{" {
					glog.Warningf("invalid lease line %q", l)
					break
				}

				ip := net.ParseIP(parts[1])
				if ip == nil {
					glog.Warningf("invalid lease line %q: invalid ip")
					break
				}
				curAddr = &ip
			}
		} else {
			parts[len(parts)-1] = strings.TrimRight(parts[len(parts)-1], ";")

			switch parts[0] {
			case "starts":
				fallthrough
			case "ends":
				if len(parts) != 4 {
					glog.Warningf("invalid time line %q", l)
					break
				}
				t, err := time.Parse(timeFmt, fmt.Sprintf("%s %s", parts[2], parts[3]))
				if err != nil {
					glog.Warningf("invalid time line %q: %v", l, err)
					break
				}
				if parts[0] == "starts" {
					curStart = &t
				} else {
					curEnd = &t
				}
			case "hardware":
				if len(parts) < 2 {
					glog.Warningf("invalid hardware line %q", l)
					break
				}
				if parts[1] != "ethernet" {
					break
				}
				if len(parts) != 3 {
					glog.Warningf("invalid hardware ethernet line %q", l)
					break
				}
				hw, err := net.ParseMAC(parts[2])
				if err != nil {
					glog.Warningf("invalid hardware ethernet line %q: %v", l, err)
					break
				}
				curHW = &hw
			case "}":
				if curStart == nil || curEnd == nil || curHW == nil {
					glog.V(2).Info("Invalid block for %q, not enough fields", curAddr.String())
				} else if curEnd.Before(now) {
					// skip.
				} else {
					leases[curHW.String()] = &lease{
						hardware: *curHW,
						ip:       *curAddr,
						from:     *curStart,
						to:       *curEnd,
					}
				}
				curAddr = nil
				curStart = nil
				curEnd = nil
				curHW = nil
			}
		}
	}

	ret := make([]lease, len(leases))
	i := 0
	for _, v := range leases {
		ret[i] = *v
		i += 1
	}

	return ret, nil
}

type service struct {
	leaseFile          string
	leaseRefreshString string
	leaseRefresh       time.Duration
	leaseC             chan chan []lease
}

func (s *service) work(ctx context.Context) {
	leases := []lease{}

	ticker := time.NewTicker(s.leaseRefresh)
	start := make(chan struct{}, 1)
	start <- struct{}{}

	work := func() {
		glog.Infof("Parsing leases...")
		f, err := os.Open(s.leaseFile)
		if err != nil {
			glog.Errorf("Could not open lease file: %v", err)
			return
		}
		l, err := parseLeases(f, time.Now())
		f.Close()
		if err != nil {
			glog.Errorf("Could not parse lease file: %v", err)
			return
		}
		sort.Slice(l, func(i, j int) bool { return string([]byte(l[i].ip)) < string([]byte(l[j].ip)) })
		leases = l
		glog.Infof("Got %d leases", len(leases))
	}

	glog.Infof("Worker started.")

	for {
		select {
		case <-start:
			work()
		case <-ticker.C:
			work()
		case c := <-s.leaseC:
			c <- leases
		case <-ctx.Done():
			glog.Infof("Worker quitting.")
			close(start)
			return
		}
	}
}

func (s *service) Leases(ctx context.Context, req *hpb.LeasifierLeasesRequest) (*hpb.LeasifierLeasesResponse, error) {
	c := make(chan []lease)
	s.leaseC <- c
	leases := <-c

	res := &hpb.LeasifierLeasesResponse{
		Leases: make([]*hpb.LeasifierLease, len(leases)),
	}

	for i, l := range leases {
		res.Leases[i] = &hpb.LeasifierLease{
			PhysicalAddress: l.hardware.String(),
			IpAddress:       l.ip.String(),
		}
	}

	return res, nil
}

func main() {
	s := &service{
		leaseC: make(chan chan []lease),
	}

	flag.StringVar(&s.leaseFile, "lease_file", "/var/db/dhcpd.leases", "Location of leasefile")
	flag.StringVar(&s.leaseRefreshString, "lease_refresh", "1m", "How often to refresh leases")
	flag.Parse()

	d, err := time.ParseDuration(s.leaseRefreshString)
	if err != nil {
		glog.Exit(err)
	}
	s.leaseRefresh = d

	m := mirko.New()
	if err := m.Listen(); err != nil {
		glog.Exitf("Could not listen: %v", err)
	}

	hpb.RegisterLeasifierServer(m.GRPC(), s)

	s.setupStatusz(m)
	go s.work(m.Context())

	if err := m.Serve(); err != nil {
		glog.Exitf("Could not run: %v", err)
	}

	glog.Info("Running.")
	<-m.Done()

}