<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - <NAME> <EMAIL> [:ulfr]
package main

import (
	"fmt"
	"io/ioutil"
	"mig.ninja/mig"
	"mig.ninja/mig/mig-agent/agentcontext"
	"net/http"
	"os"
	"time"
)

var sockCtx *Context

func initSocket(ctx *Context) {
	sockCtx = ctx
	for {
		http.HandleFunc("/pid", socketHandlePID)
		http.HandleFunc("/shutdown", socketHandleShutdown)
		err := http.ListenAndServe(ctx.Socket.Bind, nil)
		if err != nil {
			ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Error from stat socket: %q", err)}.Err()
		}
		time.Sleep(60 * time.Second)
	}
}

func socketCheckQueueloc(req *http.Request) error {
	qv := req.Header.Get("AGENTID")
	if qv == "" {
		return fmt.Errorf("must set AGENTID header")
	}
	if sockCtx.Agent.UID != qv {
		return fmt.Errorf("invalid agent id")
	}
	return nil
}

func socketHandlePID(w http.ResponseWriter, req *http.Request) {
	publication.Lock()
	defer publication.Unlock()
	fmt.Fprintf(w, "%v", os.Getpid())
}

func socketHandleShutdown(w http.ResponseWriter, req *http.Request) {
	publication.Lock()
	defer publication.Unlock()
	err := socketCheckQueueloc(req)
	if err != nil {
		http.Error(w, fmt.Sprintf("%v", err), http.StatusUnauthorized)
		return
	}
	sockCtx.Channels.Terminate <- "shutdown requested"
}
<commit_msg>[minor] add a new status endpoint to the agent status socket<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - <NAME> <EMAIL> [:ulfr]
package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"mig.ninja/mig"
	"mig.ninja/mig/mig-agent/agentcontext"
	"net/http"
	"os"
	"time"
)

var sockCtx *Context

var statusTmpl = `<html>
<head>
<style>
body {
  background-color: linen;
  font-family: "Lucida Console", Monaco, monospace;
  font-size: 14px;
}
div {
  padding-top: 12px;
  padding-bottom: 12px;
}
th, td {
  padding: 7px;
  text-align: left;
}
table, td {
  border-collapse: collapse;
  border: 1px solid black;
  font-size: 14px;
}
td:nth-child(2) {
  background-color: #eeeeee;
}
</style>
</head>
<body>
<h1>mig-agent</h1>
<div>
<table>
  <tr><th colspan=2>Agent status</th></tr>
  <tr><td>Agent name</td><td>{{.Context.Agent.Hostname}}</td></tr>
  <tr><td>BinPath</td><td>{{.Context.Agent.BinPath}}</td></tr>
  <tr><td>RunDir</td><td>{{.Context.Agent.RunDir}}</td></tr>
  <tr><td>PID</td><td>{{.Pid}}</td></tr>
  <tr><td>Environment</td><td>{{.Env}}</td></tr>
  <tr><td>Tags</td><td>{{.Tags}}</td></tr>
</table>
</div>
<div>
<table>
  <tr><th colspan=2>Configuration</th></tr>
  <tr><td>Immortal</td><td>{{.Immortal}}</td></tr>
  <tr><td>Install as a service</td><td>{{.InstallService}}</td></tr>
  <tr><td>Discover public IP</td><td>{{.DiscoverPublicIP}}</td></tr>
  <tr><td>Discover AWS metadata</td><td>{{.DiscoverAWSMeta}}</td></tr>
  <tr><td>Checkin mode</td><td>{{.Checkin}}</td></tr>
  <tr><td>Only verify pubkeys (no ACL verification)</td><td>{{.OnlyVerifyPubkey}}</td></tr>
  <tr><td>Extra privacy mode</td><td>{{.ExtraPrivacyMode}}</td></tr>
  <tr><td>Environment refresh period</td><td>{{.RefreshEnv}}</td></tr>
  <tr><td>Module configuration directory</td><td>{{.ModuleConfigDir}}</td></tr>
  <tr><td>Spawn persistent modules</td><td>{{.SpawnPersistent}}</td></tr>
  <tr><td>Proxies</td><td>{{.Proxies}}</td></tr>
  <tr><td>Heartbeat frequency</td><td>{{.HeartBeatFreq}}</td></tr>
  <tr><td>Module timeout</td><td>{{.ModuleTimeout}}</td></tr>
</table>
</div>
</body>
</html>
`

type templateData struct {
	Context *Context
	Pid     int
	Env     string
	Tags    string

	Immortal         bool
	InstallService   bool
	DiscoverPublicIP bool
	DiscoverAWSMeta  bool
	Checkin          bool
	OnlyVerifyPubkey bool
	ExtraPrivacyMode bool
	RefreshEnv       time.Duration
	ModuleConfigDir  string
	SpawnPersistent  bool
	Proxies          []string
	HeartBeatFreq    time.Duration
	ModuleTimeout    time.Duration
}

func (t *templateData) importAgentConfig() {
	t.Immortal = ISIMMORTAL
	t.InstallService = MUSTINSTALLSERVICE
	t.DiscoverPublicIP = DISCOVERPUBLICIP
	t.DiscoverAWSMeta = DISCOVERAWSMETA
	t.Checkin = CHECKIN
	t.OnlyVerifyPubkey = ONLYVERIFYPUBKEY
	t.ExtraPrivacyMode = EXTRAPRIVACYMODE
	t.RefreshEnv = REFRESHENV
	t.ModuleConfigDir = MODULECONFIGDIR
	t.SpawnPersistent = SPAWNPERSISTENT
	t.Proxies = PROXIES
	t.HeartBeatFreq = HEARTBEATFREQ
	t.ModuleTimeout = MODULETIMEOUT
}

func initSocket(ctx *Context) {
	sockCtx = ctx
	for {
		http.HandleFunc("/pid", socketHandlePID)
		http.HandleFunc("/shutdown", socketHandleShutdown)
		http.HandleFunc("/", socketHandleStatus)
		err := http.ListenAndServe(ctx.Socket.Bind, nil)
		if err != nil {
			ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Error from stat socket: %q", err)}.Err()
		}
		time.Sleep(60 * time.Second)
	}
}

func socketCheckQueueloc(req *http.Request) error {
	qv := req.Header.Get("AGENTID")
	if qv == "" {
		return fmt.Errorf("must set AGENTID header")
	}
	if sockCtx.Agent.UID != qv {
		return fmt.Errorf("invalid agent id")
	}
	return nil
}

func socketHandleStatus(w http.ResponseWriter, req *http.Request) {
	tdata := templateData{Context: sockCtx}
	tdata.Pid = os.Getpid()
	tdata.importAgentConfig()
	buf, err := json.Marshal(sockCtx.Agent.Env)
	if err != nil {
		http.Error(w, fmt.Sprintf("%v", err), http.StatusInternalServerError)
		return
	}
	tdata.Env = string(buf)
	buf, err = json.Marshal(sockCtx.Agent.Tags)
	if err != nil {
		http.Error(w, fmt.Sprintf("%v", err), http.StatusInternalServerError)
		return
	}
	tdata.Tags = string(buf)
	t, err := template.New("status").Parse(statusTmpl)
	if err != nil {
		http.Error(w, fmt.Sprintf("%v", err), http.StatusInternalServerError)
		return
	}
	err = t.Execute(w, tdata)
	if err != nil {
		fmt.Fprintf(w, "%v", err)
		return
	}
}

func socketHandlePID(w http.ResponseWriter, req *http.Request) {
	publication.Lock()
	defer publication.Unlock()
	fmt.Fprintf(w, "%v", os.Getpid())
}

func socketHandleShutdown(w http.ResponseWriter, req *http.Request) {
	publication.Lock()
	defer publication.Unlock()
	err := socketCheckQueueloc(req)
	if err != nil {
		http.Error(w, fmt.Sprintf("%v", err), http.StatusUnauthorized)
		return
	}
	sockCtx.Channels.Terminate <- "shutdown requested"
}

{'fraction_non_alphanumeric': 0.17515776512236417, 'fraction_numerical': 0.004155764198861013, 'mean_word_length': 4.3088235294117645, 'pattern_counts': {'":': 0, '<': 151, '<?xml version=': 0, '>': 147, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4374404', 'n_tokens_mistral': 2558, 'n_tokens_neox': 2442, 'n_words': 504}