main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/aws/arn"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/secretsmanager"
	secretsTypes "github.com/aws/aws-sdk-go-v2/service/secretsmanager/types"
	"github.com/aws/aws-sdk-go-v2/service/ssm"
	ssmTypes "github.com/aws/aws-sdk-go-v2/service/ssm/types"
	"go.uber.org/zap"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/clientcredentials"
	"tailscale.com/client/tailscale"
)

var logger *zap.SugaredLogger

const (
	baseURL   = "https://api.tailscale.com"
	userAgent = "tailscale-authkey-rotation-function"
)

// main initializes the logger and starts the Lambda handler.
func main() {
	z, err := zap.NewProduction()
	if err != nil {
		log.Fatalf("failed to initialize zap logger: %v", err)
	}
	defer z.Sync()

	logger = z.Sugar()
	lambda.Start(handleRequest)
}

// handleRequest is the AWS Lambda entry point.
func handleRequest(ctx context.Context) (string, error) {
	clientID := os.Getenv("TAILSCALE_CLIENT_ID")
	if clientID == "" {
		return "", fmt.Errorf("missing TAILSCALE_CLIENT_ID")
	}

	clientSecretARN := os.Getenv("TAILSCALE_CLIENT_SECRET_ARN")
	if clientSecretARN == "" {
		return "", fmt.Errorf("missing TAILSCALE_CLIENT_SECRET_ARN")
	}

	storageARN := os.Getenv("STORAGE_ARN")
	if storageARN == "" {
		return "", fmt.Errorf("missing STORAGE_ARN")
	}

	rotationPeriodStr := os.Getenv("KEY_ROTATION_PERIOD_HOURS")
	if rotationPeriodStr == "" {
		return "", fmt.Errorf("missing KEY_ROTATION_PERIOD_HOURS")
	}

	tags := os.Getenv("TAILSCALE_TAGS")
	if tags == "" {
		return "", fmt.Errorf("missing TAILSCALE_TAGS")
	}
	parsedTags := strings.Split(tags, ",")

	// Parse environment variables for ephemeral, reusable, preauthorized
	ephemeral := parseBool(os.Getenv("TAILSCALE_KEY_EPHEMERAL"), true)
	reusable := parseBool(os.Getenv("TAILSCALE_KEY_REUSABLE"), true)
	preauth := parseBool(os.Getenv("TAILSCALE_KEY_PREAUTHORIZED"), true)

	// Retrieve Tailscale client secret from SSM/Secrets Manager
	clientSecret, err := fetchClientSecret(ctx, clientSecretARN)
	if err != nil {
		return "", fmt.Errorf("failed to retrieve Tailscale client secret: %w", err)
	}

	storageType, storageName, err := parseStorageARN(storageARN)
	if err != nil {
		return "", fmt.Errorf("invalid STORAGE_ARN: %w", err)
	}

	rotationPeriod, err := time.ParseDuration(rotationPeriodStr + "h")
	if err != nil {
		return "", fmt.Errorf("failed to parse KEY_ROTATION_PERIOD_HOURS: %w", err)
	}
	// Add buffer
	rotationPeriod += time.Hour

	// Build HTTP client using the retrieved secret
	tsHTTPClient, err := getTailscaleHTTPClient(ctx, clientID, clientSecret)
	if err != nil {
		return "", fmt.Errorf("failed to obtain Tailscale OAuth token: %w", err)
	}

	newKey, err := createTailscaleKey(ctx, tsHTTPClient, rotationPeriod, ephemeral, reusable, preauth, parsedTags)
	if err != nil {
		return "", fmt.Errorf("failed to create Tailscale key: %w", err)
	}

	oldKey, err := fetchExistingKey(ctx, storageType, storageName)
	if err != nil {
		logger.Errorf("failed to fetch existing key (may not exist yet): %v", err)
	} else {
		oldKeyID, parseErr := parseTailscaleKeyID(oldKey)
		if parseErr != nil {
			logger.Infof("failed to parse old Tailscale key ID: %v", parseErr)
		} else {
			logger.Infof("old Tailscale key with ID %s found, deleting it", oldKeyID)

			delClient := tailscale.NewClient("-", nil)
			delClient.UserAgent = userAgent
			delClient.HTTPClient = tsHTTPClient
			delClient.BaseURL = baseURL

			if err := delClient.DeleteKey(ctx, oldKeyID); err != nil {
				logger.Errorf("failed to delete old Tailscale key: %v", err)
			} else {
				logger.Infof("old Tailscale key deleted successfully")
			}
		}
	}

	if err := storeNewKey(ctx, storageType, storageName, newKey); err != nil {
		return "", fmt.Errorf("failed to store new Tailscale key: %w", err)
	}

	msg := fmt.Sprintf("Successfully rotated Tailscale key in %s (resource: %s). ephemeral=%v reusable=%v preauth=%v tags=%v",
		storageType, storageName, ephemeral, reusable, preauth, parsedTags)
	logger.Info(msg)
	return msg, nil
}

// fetchClientSecret retrieves the Tailscale client secret from either SSM or Secrets Manager, depending on the ARN.
func fetchClientSecret(ctx context.Context, secretARN string) (string, error) {
	parsed, err := arn.Parse(secretARN)
	if err != nil {
		return "", fmt.Errorf("invalid TAILSCALE_CLIENT_SECRET_ARN: %w", err)
	}

	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return "", err
	}

	switch parsed.Service {
	case "ssm":
		resource := parsed.Resource
		if !strings.HasPrefix(resource, "parameter/") {
			return "", fmt.Errorf("unsupported SSM ARN format for TAILSCALE_CLIENT_SECRET_ARN: %s", resource)
		}
		paramName := strings.TrimPrefix(resource, "parameter/")
		// Optionally enforce a leading slash if multi-segment
		if strings.Contains(paramName, "/") && !strings.HasPrefix(paramName, "/") {
			paramName = "/" + paramName
		}

		ssmClient := ssm.NewFromConfig(cfg)
		withDecryption := true
		out, err := ssmClient.GetParameter(ctx, &ssm.GetParameterInput{
			Name:           aws.String(paramName),
			WithDecryption: &withDecryption,
		})
		if err != nil {
			return "", err
		}
		return aws.ToString(out.Parameter.Value), nil

	case "secretsmanager":
		smClient := secretsmanager.NewFromConfig(cfg)
		out, err := smClient.GetSecretValue(ctx, &secretsmanager.GetSecretValueInput{
			SecretId: aws.String(secretARN),
		})
		if err != nil {
			return "", err
		}
		return aws.ToString(out.SecretString), nil

	default:
		return "", fmt.Errorf("unsupported service in TAILSCALE_CLIENT_SECRET_ARN: %s", parsed.Service)
	}
}

func createTailscaleKey(
	ctx context.Context,
	tsHTTPClient *http.Client,
	expiry time.Duration,
	ephemeral, reusable, preauth bool,
	tags []string,
) (string, error) {

	tailscale.I_Acknowledge_This_API_Is_Unstable = true

	caps := tailscale.KeyCapabilities{
		Devices: tailscale.KeyDeviceCapabilities{
			Create: tailscale.KeyDeviceCreateCapabilities{
				Reusable:      reusable,
				Ephemeral:     ephemeral,
				Preauthorized: preauth,
				Tags:          tags,
			},
		},
	}

	client := tailscale.NewClient("-", nil)
	client.UserAgent = userAgent
	client.HTTPClient = tsHTTPClient
	client.BaseURL = baseURL

	authKey, _, err := client.CreateKeyWithExpiry(ctx, caps, expiry)
	if err != nil {
		return "", err
	}
	return authKey, nil
}

func parseTailscaleKeyID(key string) (string, error) {
	parts := strings.Split(key, "-")
	if len(parts) != 4 {
		return "", fmt.Errorf("invalid key format: %s", key)
	}
	return parts[2], nil
}

func parseStorageARN(arnString string) (string, string, error) {
	a, err := arn.Parse(arnString)
	if err != nil {
		return "", "", fmt.Errorf("failed to parse ARN: %w", err)
	}
	switch a.Service {
	case "ssm":
		resource := a.Resource
		if !strings.HasPrefix(resource, "parameter/") {
			return "", "", fmt.Errorf("unsupported SSM ARN format: %s", resource)
		}
		name := strings.TrimPrefix(resource, "parameter/")
		// Optionally enforce a leading slash if multi-segment
		if strings.Contains(name, "/") && !strings.HasPrefix(name, "/") {
			name = "/" + name
		}
		return "ssm", name, nil

	case "secretsmanager":
		return "secretsmanager", arnString, nil

	default:
		return "", "", fmt.Errorf("unsupported service in ARN: %s", a.Service)
	}
}

func fetchExistingKey(ctx context.Context, storageType, storageName string) (string, error) {
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return "", err
	}
	switch storageType {
	case "ssm":
		ssmClient := ssm.NewFromConfig(cfg)
		decrypt := true
		out, err := ssmClient.GetParameter(ctx, &ssm.GetParameterInput{
			Name:           aws.String(storageName),
			WithDecryption: &decrypt,
		})
		if err != nil {
			return "", err
		}
		return aws.ToString(out.Parameter.Value), nil
	case "secretsmanager":
		smClient := secretsmanager.NewFromConfig(cfg)
		out, err := smClient.GetSecretValue(ctx, &secretsmanager.GetSecretValueInput{
			SecretId: aws.String(storageName),
		})
		if err != nil {
			return "", err
		}
		return aws.ToString(out.SecretString), nil
	default:
		return "", fmt.Errorf("unsupported storage type: %s", storageType)
	}
}

func storeNewKey(ctx context.Context, storageType, storageName, newKey string) error {
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return err
	}
	switch storageType {
	case "ssm":
		ssmClient := ssm.NewFromConfig(cfg)
		overwrite := true
		_, err = ssmClient.PutParameter(ctx, &ssm.PutParameterInput{
			Name:      aws.String(storageName),
			Value:     aws.String(newKey),
			Type:      ssmTypes.ParameterTypeSecureString,
			Overwrite: &overwrite,
		})
		if err != nil {
			return err
		}
	case "secretsmanager":
		smClient := secretsmanager.NewFromConfig(cfg)
		_, err = smClient.UpdateSecret(ctx, &secretsmanager.UpdateSecretInput{
			SecretId:     aws.String(storageName),
			SecretString: aws.String(newKey),
		})
		if err != nil {
			var rnf *secretsTypes.ResourceNotFoundException
			if errors.As(err, &rnf) {
				return fmt.Errorf("secret not found for ARN=%s, cannot auto-create with same ARN: %w", storageName, err)
			}
			return err
		}
	default:
		return fmt.Errorf("unsupported storage type: %s", storageType)
	}
	return nil
}

// getTailscaleHTTPClient returns an *http.Client that uses OAuth2 client_credentials.
func getTailscaleHTTPClient(ctx context.Context, clientID, clientSecret string) (*http.Client, error) {
	cfg := clientcredentials.Config{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		TokenURL:     baseURL + "/api/v2/oauth/token",
		Scopes:       []string{},
		AuthStyle:    oauth2.AuthStyleInParams,
	}
	return cfg.Client(ctx), nil
}

func parseBool(s string, defaultVal bool) bool {
	switch strings.ToLower(s) {
	case "true", "1", "yes":
		return true
	case "false", "0", "no":
		return false
	case "":
		return defaultVal
	default:
		return defaultVal
	}
}