work life postureknyar/worklifeposture

5 files changed, 186 insertions, 0 deletions

diff --git a/cmd/worklifeposture/.gitignore b/cmd/worklifeposture/.gitignore
new file mode 100644
index 000000000..cc6618c2d
--- /dev/null
+++ b/cmd/worklifeposture/.gitignore
@@ -0,0 +1 @@
+GeoLite2-City.mmdb
diff --git a/cmd/worklifeposture/worklifeposture.go b/cmd/worklifeposture/worklifeposture.go
new file mode 100644
index 000000000..10bd211cf
--- /dev/null
+++ b/cmd/worklifeposture/worklifeposture.go
@@ -0,0 +1,168 @@
+// worklifeposture enables achieving work-life balance through device posture.
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/oschwald/geoip2-golang"
+	"golang.org/x/oauth2/clientcredentials"
+	"tailscale.com/ipn/ipnstate"
+	"tailscale.com/syncs"
+	"tailscale.com/tailcfg"
+	"tailscale.com/tsnet"
+	"tailscale.com/types/logger"
+	"tailscale.com/util/httpm"
+)
+
+func main() {
+	s := &tsnet.Server{
+		Hostname: "worklifeposture",
+		Logf:     logger.Discard,
+	}
+
+	// maxmind.com
+	db, err := geoip2.Open("GeoLite2-City.mmdb")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Close()
+
+	writer := newAttrWriter()
+
+	var lastProcessed syncs.Map[tailcfg.StableNodeID, time.Time]
+	for {
+		lc, err := s.LocalClient()
+		if err != nil {
+			log.Printf("error getting local client: %s", err)
+			continue
+		}
+		st, err := lc.Status(context.Background())
+		if err != nil {
+			log.Printf("error calling status: %s", err)
+			continue
+		}
+
+		var wg sync.WaitGroup
+		sema := make(chan struct{}, 5) // limit concurrency
+		for _, peer := range st.Peer {
+			if last, ok := lastProcessed.Load(peer.ID); ok && time.Since(last) < 5*time.Minute {
+				continue
+			}
+			sema <- struct{}{} // acquire a semaphore
+			wg.Add(1)
+			go func(peer *ipnstate.PeerStatus) {
+				defer wg.Done()
+				defer func() { <-sema }() // release the semaphore
+
+				// Ping to trigger disco.
+				ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+				_, err := lc.Ping(ctx, peer.TailscaleIPs[0], tailcfg.PingPeerAPI)
+				cancel()
+				if err != nil {
+					if !errors.Is(err, context.DeadlineExceeded) {
+						log.Printf("ping %s error: %s", peer.ID, err)
+					}
+					return
+				}
+
+				// Look up timezone based on the public IP.
+				ip := peer.CurAddr
+				if ip == "" {
+					log.Printf("no IP for peer %s", peer.ID)
+					return
+				}
+				netip := net.ParseIP(ip)
+				if netip == nil {
+					log.Printf("cannot parse IP %q for peer %s", ip, peer.ID)
+					return
+				}
+				res, err := db.City(netip)
+				if err != nil {
+					log.Printf("error looking up details for %s: %s", netip, err)
+					return
+				}
+
+				// Write an attribute depending on whether it's working hours in a given timezone.
+				tz, err := time.LoadLocation(res.Location.TimeZone)
+				if err != nil {
+					log.Printf("error loading location %s: %s", res.Location.TimeZone, err)
+					return
+				}
+				at := time.Now().In(tz)
+				value := "life"
+				if shouldYouWork(at) {
+					value = "work"
+				}
+				if err := writer.write(peer.ID, "custom:balance", value); err != nil {
+					log.Printf("error writing attribute for %s: %s", peer.ID, err)
+					return
+				}
+				lastProcessed.Store(peer.ID, time.Now())
+			}(peer)
+		}
+		wg.Wait()
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func shouldYouWork(at time.Time) bool {
+	if at.Weekday() == time.Saturday || at.Weekday() == time.Sunday {
+		return false
+	}
+	workingBegins := time.Date(at.Year(), at.Month(), at.Day(), 9, 0, 0, 0, at.Location())
+	workingFinallyEnds := time.Date(at.Year(), at.Month(), at.Day(), 17, 0, 0, 0, at.Location())
+	if at.After(workingBegins) && at.Before(workingFinallyEnds) {
+		return true
+	}
+	return false
+}
+
+type attrWriter struct {
+	client *http.Client
+}
+
+func newAttrWriter() *attrWriter {
+	var oauthConfig = &clientcredentials.Config{
+		ClientID:     os.Getenv("OAUTH_CLIENT_ID"),
+		ClientSecret: os.Getenv("OAUTH_CLIENT_SECRET"),
+		TokenURL:     "https://api.tailscale.com/api/v2/oauth/token",
+	}
+	return &attrWriter{client: oauthConfig.Client(context.Background())}
+}
+
+func (aw *attrWriter) write(nodeID tailcfg.StableNodeID, key, value string) error {
+	url := fmt.Sprintf("https://api.tailscale.com/api/v2/device/%s/attributes/%s", nodeID, key)
+
+	valueMap := map[string]any{"value": value}
+	valueBody, err := json.Marshal(valueMap)
+	if err != nil {
+		return err
+	}
+
+	setReq, err := http.NewRequest(httpm.POST, url, bytes.NewReader(valueBody))
+	if err != nil {
+		return err
+	}
+
+	resp, err := aw.client.Do(setReq)
+	if err != nil {
+		return err
+	}
+	resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		return fmt.Errorf("unexpected status for %s: %s", nodeID, resp.Status)
+	}
+	log.Printf("set %s %q=%q", nodeID, key, value)
+	return nil
+}
diff --git a/go.mod b/go.mod
index f7e1e7100..8616d209a 100644
--- a/go.mod
+++ b/go.mod
@@ -53,6 +53,7 @@ require (
 	github.com/mdlayher/sdnotify v1.0.0
 	github.com/miekg/dns v1.1.56
 	github.com/mitchellh/go-ps v1.0.0
+	github.com/oschwald/geoip2-golang v1.9.0
 	github.com/peterbourgon/ff/v3 v3.4.0
 	github.com/pkg/errors v0.9.1
 	github.com/pkg/sftp v1.13.6
@@ -115,6 +116,7 @@ require (
 	github.com/gobuffalo/flect v1.0.2 // indirect
 	github.com/google/gnostic-models v0.6.9-0.20230804172637-c7be7c783f49 // indirect
 	github.com/gorilla/securecookie v1.1.1 // indirect
+	github.com/oschwald/maxminddb-golang v1.11.0 // indirect
 )
 
 require (
diff --git a/go.sum b/go.sum
index 76f2e0e58..22a0e4b4c 100644
--- a/go.sum
+++ b/go.sum
@@ -722,6 +722,10 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0-rc5 h1:Ygwkfw9bpDvs+c9E34SdgGOj41dX/cbdlwvlWt0pnFI=
 github.com/opencontainers/image-spec v1.1.0-rc5/go.mod h1:X4pATf0uXsnn3g5aiGIsVnJBR4mxhKzfwmvK/B2NTm8=
+github.com/oschwald/geoip2-golang v1.9.0 h1:uvD3O6fXAXs+usU+UGExshpdP13GAqp4GBrzN7IgKZc=
+github.com/oschwald/geoip2-golang v1.9.0/go.mod h1:BHK6TvDyATVQhKNbQBdrj9eAvuwOMi2zSFXizL3K81Y=
+github.com/oschwald/maxminddb-golang v1.11.0 h1:aSXMqYR/EPNjGE8epgqwDay+P30hCBZIveY0WZbAWh0=
+github.com/oschwald/maxminddb-golang v1.11.0/go.mod h1:YmVI+H0zh3ySFR3w+oz8PCfglAFj3PuCmui13+P9zDg=
 github.com/otiai10/copy v1.2.0 h1:HvG945u96iNadPoG2/Ja2+AUJeW5YuFQMixq9yirC+k=
 github.com/otiai10/copy v1.2.0/go.mod h1:rrF5dJ5F0t/EWSYODDu4j9/vEeYHMkc8jt0zJChqQWw=
 github.com/otiai10/curr v0.0.0-20150429015615-9b4961190c95/go.mod h1:9qAhocn7zKJG+0mI8eUu6xqkFDYS2kb2saOteoSB3cE=
diff --git a/wgengine/magicsock/endpoint.go b/wgengine/magicsock/endpoint.go
index 58267267d..50b176688 100644
--- a/wgengine/magicsock/endpoint.go
+++ b/wgengine/magicsock/endpoint.go
@@ -1338,6 +1338,17 @@ func (de *endpoint) populatePeerStatus(ps *ipnstate.PeerStatus) {
 	if udpAddr, derpAddr, _ := de.addrForSendLocked(now); udpAddr.IsValid() && !derpAddr.IsValid() {
 		ps.CurAddr = udpAddr.String()
 	}
+
+	// If we don't have a direct connection, look for a public IP in the list of advertised endpoints.
+	if ps.CurAddr != "" {
+		return
+	}
+	for ip := range de.endpointState {
+		if ip.Addr().IsPrivate() {
+			continue
+		}
+		ps.CurAddr = ip.Addr().String()
+	}
 }
 
 // stopAndReset stops timers associated with de and resets its state back to zero.