talm

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package poweroff
6

7
import (
8
	"context"
9
	"fmt"
10
	"log"
11
	"slices"
12

13
	"google.golang.org/grpc"
14
	"google.golang.org/grpc/credentials/insecure"
15
	"google.golang.org/grpc/metadata"
16

17
	"github.com/siderolabs/talos/pkg/grpc/middleware/authz"
18
	"github.com/siderolabs/talos/pkg/machinery/client"
19
	"github.com/siderolabs/talos/pkg/machinery/constants"
20
	"github.com/siderolabs/talos/pkg/machinery/role"
21
)
22

23
// Action is the action to be performed by the poweroff command.
24
type Action string
25

26
const (
27
	// Shutdown is the action to shutdown the machine.
28
	Shutdown Action = "shutdown"
29
	// Reboot is the action to reboot the machine.
30
	Reboot Action = "reboot"
31
)
32

33
// Main is the entrypoint into /sbin/poweroff.
34
func Main(args []string) {
35
	ctx := context.Background()
36

37
	md := metadata.Pairs()
38
	authz.SetMetadata(md, role.MakeSet(role.Admin))
39
	adminCtx := metadata.NewOutgoingContext(ctx, md)
40

41
	client, err := client.New(adminCtx, client.WithUnixSocket(constants.MachineSocketPath), client.WithGRPCDialOptions(grpc.WithTransportCredentials(insecure.NewCredentials())))
42
	if err != nil {
43
		log.Fatalf(fmt.Errorf("error while creating machinery client: %w", err).Error())
44
	}
45

46
	switch ActionFromArgs(args) {
47
	case Shutdown:
48
		err = client.Shutdown(adminCtx)
49
		if err != nil {
50
			log.Fatalf(fmt.Errorf("error while sending shutdown command: %w", err).Error())
51
		}
52
	case Reboot:
53
		err = client.Reboot(adminCtx)
54
		if err != nil {
55
			log.Fatalf(fmt.Errorf("error while sending reboot command: %w", err).Error())
56
		}
57
	}
58
}
59

60
// ActionFromArgs returns the action to be performed based on the arguments.
61
func ActionFromArgs(args []string) Action {
62
	if len(args) > 1 {
63
		if slices.ContainsFunc(args[1:], func(s string) bool {
64
			return s == "--halt" || s == "-H" || s == "--poweroff" || s == "-P" || s == "-p"
65
		}) {
66
			return Shutdown
67
		}
68

69
		if slices.ContainsFunc(args[1:], func(s string) bool {
70
			return s == "--reboot" || s == "-r"
71
		}) {
72
			return Reboot
73
		}
74
	}
75

76
	return Shutdown
77
}
78