diff --git a/go.mod b/go.mod index c306403..c6e7243 100644 --- a/go.mod +++ b/go.mod @@ -31,7 +31,6 @@ require ( ) require ( - github.com/ahmetb/go-linq/v3 v3.2.0 github.com/aws/aws-sdk-go-v2/credentials v1.20.9 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.20.3 // indirect github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.14.6 // indirect diff --git a/go.sum b/go.sum index 76c72f0..f69da1c 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,3 @@ -github.com/ahmetb/go-linq/v3 v3.2.0 h1:BEuMfp+b59io8g5wYzNoFe9pWPalRklhlhbiU3hYZDE= -github.com/ahmetb/go-linq/v3 v3.2.0/go.mod h1:haQ3JfOeWK8HpVxMtHHEMPVgBKiYyQ+f1/kLZh/cj9U= github.com/apex/log v1.9.0 h1:FHtw/xuaM8AgmvDDTI9fiwoAL25Sq2cxojnZICUU8l0= github.com/apex/log v1.9.0/go.mod h1:m82fZlWIuiWzWP04XCTXmnX0xRkYYbCdYn8jbJeLBEA= github.com/apex/logs v1.0.0/go.mod h1:XzxuLZ5myVHDy9SAmYpamKKRNApGj54PfYLcFrXqDwo= @@ -217,8 +215,6 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= -gopkg.in/ini.v1 v1.67.1 h1:tVBILHy0R6e4wkYOn3XmiITt/hEVH4TFMYvAX2Ytz6k= -gopkg.in/ini.v1 v1.67.1/go.mod h1:x/cyOwCgZqOkJoDIJ3c1KNHMo10+nLGAhh+kn3Zizss= gopkg.in/ini.v1 v1.67.3 h1:iM9Lhz5MRSGhHVGGwCuzG9KO8PoirCXj/m/qTmOJJQw= gopkg.in/ini.v1 v1.67.3/go.mod h1:x/cyOwCgZqOkJoDIJ3c1KNHMo10+nLGAhh+kn3Zizss= gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= diff --git a/lib/aws.go b/lib/aws.go index 5888026..0d7100d 100644 --- a/lib/aws.go +++ b/lib/aws.go @@ -5,7 +5,6 @@ import ( "fmt" "sort" - linq "github.com/ahmetb/go-linq/v3" "github.com/apex/log" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/config" @@ -58,118 +57,100 @@ func TraverseProfiles(profiles []ProfileConfig, noProfilePrefix bool) ([]Process sort.Slice(profileSummaries, func(i, j int) bool { return profileSummaries[i].Name < profileSummaries[j].Name }) var processedProfileSummaries []ProcessedProfileSummary - // go through all profileSummaries and - // create sshEntries out of it for _, summary := range profileSummaries { - var profileSSHEntries []SSHEntry - - ctx := log.WithField("profile", summary.Name) - // group instances by VPC - ctx.Debug("Grouping instances by VPC") - - var vpcInstances []linq.Group - - // take the instances slice - linq.From(summary.Instances).OrderBy(instanceNameSorter). // sort by name first - ThenBy(instanceLaunchTimeSorter). // then by launch time - GroupBy(func(i interface{}) interface{} { // and then group by vpc - vpcID := i.(types.Instance).VpcId - return aws.ToString(vpcID) - }, func(i interface{}) interface{} { - return i.(types.Instance) - }).ToSlice(&vpcInstances) - - var commonBastions []types.Instance - linq.From(summary.Instances).OrderBy(instanceNameSorter). // sort by name first - ThenBy(instanceLaunchTimeSorter). // then by launch time - Where( - func(f interface{}) bool { - return isBastionFromTags(f.(types.Instance).Tags, true) // check for global tag as well - }, - ).ToSlice(&commonBastions) - - ctx.Debugf("Found %d common (global) bastions", len(commonBastions)) - - for _, vpcGroup := range vpcInstances { // take the instances grouped by vpc and iterate - var vpcBastions []types.Instance - linq.From(vpcGroup.Group).Where( - func(f interface{}) bool { - return isBastionFromTags(f.(types.Instance).Tags, false) // "false" means don't check for global tag - }, - ).ToSlice(&vpcBastions) - - ctx.WithField("vpc", vpcGroup.Key).Debugf("Found %d bastions", len(vpcBastions)) - - var nameInstances []linq.Group - linq.From(vpcGroup.Group).GroupBy(func(i interface{}) interface{} { // now group them by name - instanceName := getNameFromTags(i.(types.Instance).Tags) - return instanceName - }, func(i interface{}) interface{} { - return i.(types.Instance) - }).ToSlice(&nameInstances) - - // now we have instances, grouped by vpc and name - for _, nameGroup := range nameInstances { - instanceName := nameGroup.Key.(string) - - for n, instance := range nameGroup.Group { - instance := instance.(types.Instance) - var entry = SSHEntry{ - InstanceID: aws.ToString(instance.InstanceId), - ProfileConfig: ProfileConfig{ - Name: summary.Name, - Region: summary.Region, - Domain: summary.Domain, - }, - } - entry.User = GetUserFromTags(instance.Tags) - entry.Port = getPortFromTags(instance.Tags) + processedProfileSummaries = append(processedProfileSummaries, processProfileSummary(summary, noProfilePrefix)) + } + return processedProfileSummaries, errors +} - // first try to find a bastion from this vpc - bastion := findBestBastion(instanceName, vpcBastions) - if bastion == nil { // then try common ones - bastion = findBestBastion(instanceName, commonBastions) - } - entry.Address = aws.ToString(instance.PrivateIpAddress) // get the private address first as we always have one - if bastion != nil { // get private address and add proxyhost, which is the bastion ip - // refer to the bastion by its instance ID - // which we should have a record for - entry.ProxyJump = aws.ToString(bastion.InstanceId) - } else { // get public IP if we have one - if publicIP := aws.ToString(instance.PublicIpAddress); publicIP != "" { - entry.Address = aws.ToString(instance.PublicIpAddress) - } - } - var instanceIndex string - if len(nameGroup.Group) > 1 { - instanceIndex = fmt.Sprintf("%d", n+1) - } - // add all names of the instance - var name = getInstanceCanonicalName(summary.Name, instanceName, instanceIndex) - if noProfilePrefix { - name = getInstanceCanonicalName("", instanceName, instanceIndex) - } - entry.Names = append(entry.Names, name, entry.InstanceID, fmt.Sprintf("%s.%s", entry.Address, entry.ProfileConfig.Name)) - if summary.Domain != "" { - entry.Names = append(entry.Names, fmt.Sprintf("%s.%s", name, summary.Domain)) +// processProfileSummary turns raw instances of a profile into ssh entries +func processProfileSummary(summary profileSummary, noProfilePrefix bool) ProcessedProfileSummary { + var profileSSHEntries []SSHEntry + + ctx := log.WithField("profile", summary.Name) + // group instances by VPC + ctx.Debug("Grouping instances by VPC") + + // sort by name first, then by launch time + sorted := sortInstances(summary.Instances) + + vpcGroups := groupInstances(sorted, func(i types.Instance) string { return aws.ToString(i.VpcId) }) + + commonBastions := filterInstances(sorted, func(i types.Instance) bool { + return isBastionFromTags(i.Tags, true) // check for global tag as well + }) + + ctx.Debugf("Found %d common (global) bastions", len(commonBastions)) + + for _, vpcGroup := range vpcGroups { // take the instances grouped by vpc and iterate + vpcBastions := filterInstances(vpcGroup.Instances, func(i types.Instance) bool { + return isBastionFromTags(i.Tags, false) // "false" means don't check for global tag + }) + + ctx.WithField("vpc", vpcGroup.Key).Debugf("Found %d bastions", len(vpcBastions)) + + // now group them by name + nameGroups := groupInstances(vpcGroup.Instances, func(i types.Instance) string { return getNameFromTags(i.Tags) }) + + // now we have instances, grouped by vpc and name + for _, nameGroup := range nameGroups { + instanceName := nameGroup.Key + + for n, instance := range nameGroup.Instances { + var entry = SSHEntry{ + InstanceID: aws.ToString(instance.InstanceId), + ProfileConfig: ProfileConfig{ + Name: summary.Name, + Region: summary.Region, + Domain: summary.Domain, + }, + } + entry.User = GetUserFromTags(instance.Tags) + entry.Port = getPortFromTags(instance.Tags) + + // first try to find a bastion from this vpc + bastion := findBestBastion(instanceName, vpcBastions) + if bastion == nil { // then try common ones + bastion = findBestBastion(instanceName, commonBastions) + } + entry.Address = aws.ToString(instance.PrivateIpAddress) // get the private address first as we always have one + if bastion != nil { // get private address and add proxyhost, which is the bastion ip + // refer to the bastion by its instance ID + // which we should have a record for + entry.ProxyJump = aws.ToString(bastion.InstanceId) + } else { // get public IP if we have one + if publicIP := aws.ToString(instance.PublicIpAddress); publicIP != "" { + entry.Address = aws.ToString(instance.PublicIpAddress) } - profileSSHEntries = append(profileSSHEntries, entry) } + var instanceIndex string + if len(nameGroup.Instances) > 1 { + instanceIndex = fmt.Sprintf("%d", n+1) + } + // add all names of the instance + var name = getInstanceCanonicalName(summary.Name, instanceName, instanceIndex) + if noProfilePrefix { + name = getInstanceCanonicalName("", instanceName, instanceIndex) + } + entry.Names = append(entry.Names, name, entry.InstanceID, fmt.Sprintf("%s.%s", entry.Address, entry.ProfileConfig.Name)) + if summary.Domain != "" { + entry.Names = append(entry.Names, fmt.Sprintf("%s.%s", name, summary.Domain)) + } + profileSSHEntries = append(profileSSHEntries, entry) } } - // sort by the first (main) name alphabetically - sort.SliceStable(profileSSHEntries, func(i, j int) bool { return profileSSHEntries[i].Names[0] < profileSSHEntries[j].Names[0] }) - - processedProfileSummaries = append(processedProfileSummaries, ProcessedProfileSummary{ - ProfileConfig: ProfileConfig{ - Name: summary.Name, - Region: summary.Region, - Domain: summary.Domain, - }, - SSHEntries: profileSSHEntries, - }) } - return processedProfileSummaries, errors + // sort by the first (main) name alphabetically + sort.SliceStable(profileSSHEntries, func(i, j int) bool { return profileSSHEntries[i].Names[0] < profileSSHEntries[j].Names[0] }) + + return ProcessedProfileSummary{ + ProfileConfig: ProfileConfig{ + Name: summary.Name, + Region: summary.Region, + Domain: summary.Domain, + }, + SSHEntries: profileSSHEntries, + } } // DescribeProfile describes the specified profile diff --git a/lib/aws_test.go b/lib/aws_test.go new file mode 100644 index 0000000..56ac1b5 --- /dev/null +++ b/lib/aws_test.go @@ -0,0 +1,103 @@ +package lib + +import ( + "reflect" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/ec2/types" +) + +func tag(k, v string) types.Tag { return types.Tag{Key: aws.String(k), Value: aws.String(v)} } + +func inst(id, vpc, name, private, public string, launched time.Time, extra ...types.Tag) types.Instance { + i := types.Instance{ + InstanceId: aws.String(id), + VpcId: aws.String(vpc), + PrivateIpAddress: aws.String(private), + LaunchTime: aws.Time(launched), + Tags: append([]types.Tag{tag("Name", name)}, extra...), + } + if public != "" { + i.PublicIpAddress = aws.String(public) + } + return i +} + +func testSummary() profileSummary { + early := time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC) + late := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + + return profileSummary{ + ProfileConfig: ProfileConfig{Name: "prod", Region: "ap-southeast-2", Domain: "example.com"}, + // deliberately unordered + Instances: []types.Instance{ + inst("i-web2", "vpc-a", "web", "10.0.0.2", "", late), + inst("i-worker", "vpc-c", "worker", "10.2.0.1", "", early), + inst("i-shared", "vpc-b", "shared-bastion", "10.1.0.9", "9.9.9.9", early, tag("Global", "yes")), + inst("i-db", "vpc-a", "db", "10.0.0.3", "1.1.1.1", early, tag("x-aws-ssh-port", "2222")), + inst("i-api", "vpc-b", "api", "10.1.0.1", "", early), + inst("i-web1", "vpc-a", "web", "10.0.0.1", "", early, tag("x-aws-ssh-user", "ec2-user")), + inst("i-bast-a", "vpc-a", "prod-bastion", "10.0.0.9", "8.8.8.8", early), + }, + } +} + +func entry(id, addr, jump, port, user string, names ...string) SSHEntry { + return SSHEntry{ + ProfileConfig: ProfileConfig{Name: "prod", Region: "ap-southeast-2", Domain: "example.com"}, + InstanceID: id, + Address: addr, + ProxyJump: jump, + Port: port, + User: user, + Names: names, + } +} + +func TestProcessProfileSummary(t *testing.T) { + tests := []struct { + name string + noProfilePrefix bool + want []SSHEntry + }{ + { + name: "with profile prefix", + want: []SSHEntry{ + entry("i-api", "10.1.0.1", "i-shared", "", "", "prod-api", "i-api", "10.1.0.1.prod", "prod-api.example.com"), + entry("i-bast-a", "8.8.8.8", "", "", "", "prod-bastion", "i-bast-a", "8.8.8.8.prod", "prod-bastion.example.com"), + entry("i-db", "10.0.0.3", "i-bast-a", "2222", "", "prod-db", "i-db", "10.0.0.3.prod", "prod-db.example.com"), + entry("i-shared", "9.9.9.9", "", "", "", "prod-shared-bastion", "i-shared", "9.9.9.9.prod", "prod-shared-bastion.example.com"), + entry("i-web1", "10.0.0.1", "i-bast-a", "", "ec2-user", "prod-web-1", "i-web1", "10.0.0.1.prod", "prod-web-1.example.com"), + entry("i-web2", "10.0.0.2", "i-bast-a", "", "", "prod-web-2", "i-web2", "10.0.0.2.prod", "prod-web-2.example.com"), + entry("i-worker", "10.2.0.1", "i-shared", "", "", "prod-worker", "i-worker", "10.2.0.1.prod", "prod-worker.example.com"), + }, + }, + { + name: "without profile prefix", + noProfilePrefix: true, + want: []SSHEntry{ + entry("i-api", "10.1.0.1", "i-shared", "", "", "api", "i-api", "10.1.0.1.prod", "api.example.com"), + entry("i-db", "10.0.0.3", "i-bast-a", "2222", "", "db", "i-db", "10.0.0.3.prod", "db.example.com"), + entry("i-bast-a", "8.8.8.8", "", "", "", "prod-bastion", "i-bast-a", "8.8.8.8.prod", "prod-bastion.example.com"), + entry("i-shared", "9.9.9.9", "", "", "", "shared-bastion", "i-shared", "9.9.9.9.prod", "shared-bastion.example.com"), + entry("i-web1", "10.0.0.1", "i-bast-a", "", "ec2-user", "web-1", "i-web1", "10.0.0.1.prod", "web-1.example.com"), + entry("i-web2", "10.0.0.2", "i-bast-a", "", "", "web-2", "i-web2", "10.0.0.2.prod", "web-2.example.com"), + entry("i-worker", "10.2.0.1", "i-shared", "", "", "worker", "i-worker", "10.2.0.1.prod", "worker.example.com"), + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // run repeatedly, group order must not leak into the result + for n := 0; n < 20; n++ { + got := processProfileSummary(testSummary(), tt.noProfilePrefix) + if !reflect.DeepEqual(got.SSHEntries, tt.want) { + t.Fatalf("run %d:\n got: %+v\nwant: %+v", n, got.SSHEntries, tt.want) + } + } + }) + } +} diff --git a/lib/util.go b/lib/util.go index 91314e6..f1d34a1 100644 --- a/lib/util.go +++ b/lib/util.go @@ -1,7 +1,9 @@ package lib import ( + "cmp" "regexp" + "slices" "sort" "strings" @@ -121,12 +123,47 @@ func getInstanceCanonicalName(profile, instanceName, instanceIndex string) strin return sanitiser.ReplaceAllString(strings.Join(parts, "-"), "-") } -func instanceLaunchTimeSorter(i interface{}) interface{} { // sorts by launch time - launched := aws.ToTime(i.(types.Instance).LaunchTime) - return launched.Unix() +// instanceGroup is a set of instances sharing the same key +type instanceGroup struct { + Key string + Instances []types.Instance } -func instanceNameSorter(i interface{}) interface{} { // sort by instance name - instanceName := getNameFromTags(i.(types.Instance).Tags) - return instanceName +// sortInstances returns a copy sorted by name, then by launch time +func sortInstances(instances []types.Instance) []types.Instance { + sorted := slices.Clone(instances) + slices.SortStableFunc(sorted, func(a, b types.Instance) int { + return cmp.Or( + cmp.Compare(getNameFromTags(a.Tags), getNameFromTags(b.Tags)), + cmp.Compare(aws.ToTime(a.LaunchTime).Unix(), aws.ToTime(b.LaunchTime).Unix()), + ) + }) + return sorted +} + +// groupInstances groups instances by key, keeping the order in which keys first appear +func groupInstances(instances []types.Instance, key func(types.Instance) string) []instanceGroup { + var groups []instanceGroup + position := make(map[string]int) + for _, instance := range instances { + k := key(instance) + n, ok := position[k] + if !ok { + n = len(groups) + position[k] = n + groups = append(groups, instanceGroup{Key: k}) + } + groups[n].Instances = append(groups[n].Instances, instance) + } + return groups +} + +func filterInstances(instances []types.Instance, keep func(types.Instance) bool) []types.Instance { + var kept []types.Instance + for _, instance := range instances { + if keep(instance) { + kept = append(kept, instance) + } + } + return kept } diff --git a/lib/util_test.go b/lib/util_test.go new file mode 100644 index 0000000..bb578ff --- /dev/null +++ b/lib/util_test.go @@ -0,0 +1,76 @@ +package lib + +import ( + "reflect" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/ec2/types" +) + +func ids(instances []types.Instance) []string { + var out []string + for _, i := range instances { + out = append(out, aws.ToString(i.InstanceId)) + } + return out +} + +func TestSortInstances(t *testing.T) { + early := time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC) + late := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + input := []types.Instance{ + inst("i-b-late", "v", "b", "1", "", late), + inst("i-a-tie1", "v", "a", "1", "", early), + inst("i-b-early", "v", "B", "1", "", early), + inst("i-a-tie2", "v", "a", "1", "", early), + } + + got := ids(sortInstances(input)) + want := []string{"i-a-tie1", "i-a-tie2", "i-b-early", "i-b-late"} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %v, want %v", got, want) + } + if first := aws.ToString(input[0].InstanceId); first != "i-b-late" { + t.Errorf("input was modified, first is now %s", first) + } +} + +func TestGroupInstances(t *testing.T) { + early := time.Time{} + input := []types.Instance{ + inst("i-1", "vpc-b", "x", "1", "", early), + inst("i-2", "vpc-a", "x", "1", "", early), + inst("i-3", "vpc-b", "x", "1", "", early), + } + + groups := groupInstances(input, func(i types.Instance) string { return aws.ToString(i.VpcId) }) + + var keys []string + for _, g := range groups { + keys = append(keys, g.Key) + } + if want := []string{"vpc-b", "vpc-a"}; !reflect.DeepEqual(keys, want) { + t.Errorf("group keys: got %v, want %v (first-seen order)", keys, want) + } + if got, want := ids(groups[0].Instances), []string{"i-1", "i-3"}; !reflect.DeepEqual(got, want) { + t.Errorf("first group: got %v, want %v", got, want) + } +} + +func TestFilterInstances(t *testing.T) { + input := []types.Instance{ + inst("i-1", "v", "keep", "1", "", time.Time{}), + inst("i-2", "v", "drop", "1", "", time.Time{}), + inst("i-3", "v", "keep", "1", "", time.Time{}), + } + + got := ids(filterInstances(input, func(i types.Instance) bool { return getNameFromTags(i.Tags) == "keep" })) + if want := []string{"i-1", "i-3"}; !reflect.DeepEqual(got, want) { + t.Errorf("got %v, want %v", got, want) + } + if got := filterInstances(input, func(types.Instance) bool { return false }); len(got) != 0 { + t.Errorf("expected nothing, got %v", ids(got)) + } +}