From 72e76cdd9ad20299e2229c0d8d872ed4e539776b Mon Sep 17 00:00:00 2001
From: Christian Muehlhaeuser <muesli@gmail.com>
Date: Sun, 11 Mar 2018 01:43:21 +0100
Subject: [PATCH] Basic unit creation functionality

---
 createcmd.go | 86 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 targets.go   | 40 ++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 3 deletions(-)
 create mode 100644 targets.go

diff --git a/createcmd.go b/createcmd.go
index 2e29e86..a901939 100644
--- a/createcmd.go
+++ b/createcmd.go
@@ -2,8 +2,12 @@ package main
 
 import (
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
+	"strconv"
 
+	"github.com/coreos/go-systemd/unit"
 	"github.com/spf13/cobra"
 )
 
@@ -15,7 +19,7 @@ type CreateOptions struct {
 	User             string
 	Group            string
 
-	RestartSec uint
+	RestartSec uint64
 	Restart    string
 
 	After    string
@@ -26,11 +30,14 @@ var (
 	createOpts = CreateOptions{}
 
 	createCmd = &cobra.Command{
-		Use:   "create <executable> <description>",
+		Use:   "create <executable> <description> <after> <wanted-by>",
 		Short: "creates a new Unit file",
 		Long:  `The create command creates a new systemd Unit file`,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			var err error
+			ts, err := targets()
+			if err != nil {
+				return fmt.Errorf("Can't find systemd targets: %s", err)
+			}
 
 			if len(args) >= 1 {
 				createOpts.Executable = args[0]
@@ -61,17 +68,90 @@ var (
 				return fmt.Errorf("create needs a description for this service")
 			}
 
+			if len(args) >= 4 {
+				createOpts.WantedBy = args[3]
+			}
+			if len(args) >= 3 {
+				createOpts.After = args[2]
+			}
+			if len(createOpts.After) == 0 || len(createOpts.WantedBy) == 0 {
+				fmt.Println("Available targets:")
+				for _, t := range ts {
+					fmt.Printf("%s - %s\n", t.Name, t.Description)
+				}
+			}
+
+			if len(createOpts.After) == 0 {
+				createOpts.After, _ = readString("Start after target", true)
+				if len(createOpts.After) == 0 {
+					return fmt.Errorf("create needs a target after which this service will be started")
+				}
+			}
+			if !ts.Contains(createOpts.After) {
+				return fmt.Errorf("Can't create service: no such target")
+			}
+
+			if len(createOpts.WantedBy) == 0 {
+				createOpts.WantedBy, _ = readString("Which target should this service be wanted by", true)
+				if len(createOpts.WantedBy) == 0 {
+					return fmt.Errorf("create needs a target which this service will be wanted by")
+				}
+			}
+			if !ts.Contains(createOpts.WantedBy) {
+				return fmt.Errorf("Can't create service: no such target")
+			}
+
 			return executeCreate()
 		},
 	}
 )
 
 func executeCreate() error {
+	u := []*unit.UnitOption{
+		&unit.UnitOption{"Unit", "Description", createOpts.Description},
+		&unit.UnitOption{"Unit", "After", createOpts.After},
+
+		&unit.UnitOption{"Service", "ExecStart", createOpts.Executable},
+		&unit.UnitOption{"Service", "User", createOpts.User},
+		&unit.UnitOption{"Service", "Group", createOpts.Group},
+		&unit.UnitOption{"Service", "Restart", createOpts.Restart},
+		&unit.UnitOption{"Service", "RestartSec", strconv.FormatUint(createOpts.RestartSec, 10)},
+
+		&unit.UnitOption{"Install", "WantedBy", createOpts.WantedBy},
+	}
+
+	r := unit.Serialize(u)
+	b, err := ioutil.ReadAll(r)
+	if err != nil {
+		return fmt.Errorf("encountered error while reading output: %v", err)
+	}
+
+	filename := filepath.Base(createOpts.Executable) + ".service"
+	f, err := os.Create(filename)
+	if err != nil {
+		return fmt.Errorf("Can't write file: %s", err)
+	}
+	defer f.Close()
+
+	_, err = f.Write(b)
+	if err != nil {
+		return fmt.Errorf("Can't write to file: %s", err)
+	}
+
+	fmt.Printf("Generated Unit file: %s\n%s\n", filename, b)
 	return nil
 }
 
 func init() {
 	createCmd.PersistentFlags().StringVarP(&createOpts.WorkingDirectory, "workingdir", "w", "", "WorkingDirectory of the service")
+	createCmd.PersistentFlags().StringVarP(&createOpts.User, "user", "u", "root", "User to run service as")
+	createCmd.PersistentFlags().StringVarP(&createOpts.Group, "group", "g", "root", "Group to run service as")
+
+	createCmd.PersistentFlags().StringVarP(&createOpts.Restart, "restart", "r", "on-failure", "When to restart the service")
+	createCmd.PersistentFlags().Uint64VarP(&createOpts.RestartSec, "restartsec", "s", 5, "How many seconds between restarts")
+
+	createCmd.PersistentFlags().StringVarP(&createOpts.After, "after", "a", "", "Target after which the service will be started")
+	createCmd.PersistentFlags().StringVarP(&createOpts.WantedBy, "wantedby", "b", "", "This service is wanted by this target")
 
 	RootCmd.AddCommand(createCmd)
 }
diff --git a/targets.go b/targets.go
new file mode 100644
index 0000000..3927809
--- /dev/null
+++ b/targets.go
@@ -0,0 +1,40 @@
+package main
+
+import (
+	"strings"
+
+	"github.com/coreos/go-systemd/dbus"
+)
+
+type Targets []dbus.UnitStatus
+
+func targets() (Targets, error) {
+	res := []dbus.UnitStatus{}
+	conn, err := dbus.New()
+	if err != nil {
+		return res, err
+	}
+	us, err := conn.ListUnits()
+	if err != nil {
+		return res, err
+	}
+	for _, v := range us {
+		if !strings.HasSuffix(v.Name, ".target") {
+			continue
+		}
+
+		res = append(res, v)
+	}
+
+	return res, nil
+}
+
+func (ts Targets) Contains(name string) bool {
+	for _, t := range ts {
+		if t.Name == name {
+			return true
+		}
+	}
+
+	return false
+}