main.go

package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"regexp"
	"syscall"

	"github.com/aws/aws-sdk-go/service/ec2instanceconnect"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	instanceIdentifier = kingpin.Flag("instance-identifier", "Identifier for the instance").Required().String()
	osUser             = kingpin.Flag("os-user", "OS user to use").Default("centos").String()
	publicKey          = kingpin.Flag("public-key", "Public key file to use").Default("/home/username/.ssh/id_rsa.pub").String()
)

func FindInstance(sess *session.Session, instanceIdentifier string) (*ec2.Instance, error) {
	svc := ec2.New(sess)

	var err error
	var result *ec2.DescribeInstancesOutput

	isInstanceID, _ := regexp.Match("^i-[0-9a-f]{8,32}$", []byte(instanceIdentifier))

	if isInstanceID {
		result, err = svc.DescribeInstances(&ec2.DescribeInstancesInput{
			InstanceIds: []*string{
				aws.String(instanceIdentifier),
			},
		})
	} else if net.ParseIP(instanceIdentifier) != nil {
		result, err = svc.DescribeInstances(&ec2.DescribeInstancesInput{
			Filters: []*ec2.Filter{
				{
					Name: aws.String("private-ip-address"),
					Values: []*string{
						aws.String(instanceIdentifier),
					},
				},
			},
		})

	} else {
		result, err = svc.DescribeInstances(&ec2.DescribeInstancesInput{
			Filters: []*ec2.Filter{
				{
					Name: aws.String("tag:Name"),
					Values: []*string{
						aws.String(instanceIdentifier),
					},
				},
			},
		})
	}

	if err != nil {
		return nil, err
	}

	var foundInstance *ec2.Instance
	for _, reservation := range result.Reservations {
		for _, instance := range reservation.Instances {
			foundInstance = instance
		}
	}

	if foundInstance == nil {
		return nil, errors.New(fmt.Sprintf("no instance found with identifier: %s", instanceIdentifier))
	}

	return foundInstance, nil
}

func main() {
	kingpin.Parse()

	session, err := session.NewSession()
	if err != nil {
		log.Fatal(err)
	}

	instance, err := FindInstance(session, *instanceIdentifier)
	if err != nil {
		log.Fatal(err)
	}

	publicKey, err := ioutil.ReadFile(*publicKey)
	if err != nil {
		log.Fatal(err)
	}

	svc := ec2instanceconnect.New(session)
	input := &ec2instanceconnect.SendSSHPublicKeyInput{
		AvailabilityZone: instance.Placement.AvailabilityZone,
		InstanceId:       instance.InstanceId,
		InstanceOSUser:   osUser,
		SSHPublicKey:     aws.String(string(publicKey)),
	}
	_, err = svc.SendSSHPublicKey(input)

	if err != nil {
		log.Fatal(err)
	}

	binary, err := exec.LookPath("aws")
	if err != nil {
		log.Fatal(err)
	}

	args := []string{
		"aws", "ssm", "start-session", fmt.Sprintf("--target=%s", *instance.InstanceId), "--document-name=AWS-StartSSHSession",
	}

	execErr := syscall.Exec(binary, args, os.Environ())
	if err != nil {
		log.Fatal(execErr)
	}
}

ssh_config

Host *
  User ec2-user
  ProxyCommand sh -c "AWS_REGION=eu-central-1 /home/username/work/sqlssh/cli --instance-identifier=%h --os-user=ec2-user"