From 8defb738fcadbd028bcedfd46f0ef7b83418a518 Mon Sep 17 00:00:00 2001
From: "J. Lowell Wofford" <lowell@lanl.gov>
Date: Mon, 30 Nov 2020 10:59:38 -0700
Subject: [PATCH] restructure entropy command; streamline pkg; add addto
 function

---
 cmd/entropy/entropy.go       | 69 ++++++++++++++++++++++++++++++++----
 pkg/entropy/entropy.go       |  8 +++--
 pkg/entropy/entropy_linux.go | 17 ++++++---
 3 files changed, 81 insertions(+), 13 deletions(-)

diff --git a/cmd/entropy/entropy.go b/cmd/entropy/entropy.go
index a9ec452..e27b4d6 100644
--- a/cmd/entropy/entropy.go
+++ b/cmd/entropy/entropy.go
@@ -1,15 +1,70 @@
 package main
 
 import (
-	"log"
+	"fmt"
+	"os"
+	"strconv"
 
 	"github.com/jlowellwofford/entropy/pkg/entropy"
 )
 
-func main() {
-	ent, err := entropy.GetEntropy()
-	if err != nil {
-		log.Fatalf("failed to get entropy: %v", err)
-	}
-	log.Printf("current entropy: %d", ent)
+func usage() {
+	fmt.Printf(`
+Usage: %s <command> [<opts>...]
+
+Commands:
+
+	get(entropy)        - get the current system entropy.
+	addto(entcnt) <num> - Add <num> to the current entropy count (must be root).
+
+`, os.Args[0])
+}
+
+func usageFatal(str string, args ...interface{}) {
+	fmt.Printf("\n"+str+"\n", args...)
+	usage()
+	os.Exit(1)
+}
+
+func fatal(str string, args ...interface{}) {
+	fmt.Printf(str+"\n", args...)
+	os.Exit(1)
+}
+
+func main() {
+	var err error
+	if len(os.Args) < 2 {
+		usage()
+		os.Exit(1)
+	}
+
+	switch os.Args[1] {
+	case "get":
+		fallthrough
+	case "getentropy":
+		var ent int
+		if ent, err = entropy.GetEntCnt(); err != nil {
+			fatal("failed to get entropy: %v", err)
+		}
+		fmt.Printf("%d\n", ent)
+	case "addto":
+		fallthrough
+	case "addtoentcnt":
+		if len(os.Args) != 3 {
+			usageFatal("addtoentcnt requires a number to add as an option")
+		}
+		var add int
+		if add, err = strconv.Atoi(os.Args[2]); err != nil {
+			usageFatal("provided option does not appear to be a valid number: %v", err)
+		}
+		if err = entropy.AddToEntCnt(add); err != nil {
+			fatal("entropy addition failed: %v", err)
+		}
+	case "help":
+		fallthrough
+	case "usage":
+		usage()
+	default:
+		usageFatal("unrecognized command: %s", os.Args[1])
+	}
 }
diff --git a/pkg/entropy/entropy.go b/pkg/entropy/entropy.go
index a222b89..02c0d35 100644
--- a/pkg/entropy/entropy.go
+++ b/pkg/entropy/entropy.go
@@ -1,5 +1,9 @@
 package entropy
 
-func GetEntropy() (int, error) {
-	return getEntropy()
+func GetEntCnt() (int, error) {
+	return getEntCnt()
+}
+
+func AddToEntCnt(add int) error {
+	return addToEntCnt(add)
 }
diff --git a/pkg/entropy/entropy_linux.go b/pkg/entropy/entropy_linux.go
index cd17911..15df4a5 100644
--- a/pkg/entropy/entropy_linux.go
+++ b/pkg/entropy/entropy_linux.go
@@ -10,22 +10,31 @@ import (
 
 var entropy_device = "/dev/urandom"
 
-// this is honestly easier through /proc, but in the spirit of completeness...
-func getEntropy() (ent int, err error) {
+func entropyIoctl(request int, data uintptr) (err error) {
 	var fd int
 	if fd, err = unix.Open(entropy_device, unix.O_RDWR, 0); err != nil {
-		return ent, fmt.Errorf("could not open entropy device (%s): %v", entropy_device, err)
+		return fmt.Errorf("could not open entropy device (%s): %v", entropy_device, err)
 	}
 	defer unix.Close(fd)
 
 	var errno syscall.Errno
-	_, _, errno = unix.Syscall(unix.SYS_IOCTL, uintptr(fd), uintptr(RNDGETENTCNT), uintptr(unsafe.Pointer(&ent)))
+	_, _, errno = unix.Syscall(unix.SYS_IOCTL, uintptr(fd), uintptr(request), data)
 	if errno != 0 {
 		err = errno
 	}
+	return err
+}
+
+// this is honestly easier through /proc, but in the spirit of completeness...
+func getEntCnt() (ent int, err error) {
+	err = entropyIoctl(RNDGETENTCNT, uintptr(unsafe.Pointer(&ent)))
 	return
 }
 
+func addToEntCnt(add int) (err error) {
+	return entropyIoctl(RNDADDTOENTCNT, uintptr(unsafe.Pointer(&add)))
+}
+
 func addToEntropy() {
 
 }