/
main.go
122 lines (106 loc) · 2.66 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package main
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"fmt"
"os"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
_ "github.com/sheacloud/cloud-inventory/docs"
"github.com/sirupsen/logrus"
"github.com/spf13/viper"
)
var (
logrusLevels = map[string]logrus.Level{
"debug": logrus.DebugLevel,
"info": logrus.InfoLevel,
"warn": logrus.WarnLevel,
"error": logrus.ErrorLevel,
"fatal": logrus.FatalLevel,
}
)
func initOptions() {
viper.SetEnvPrefix("cloud_inventory")
viper.AutomaticEnv()
viper.BindEnv("log_level")
viper.SetDefault("log_level", "info")
viper.BindEnv("log_caller")
viper.SetDefault("log_caller", false)
viper.BindEnv("api_keys_table")
}
func initLogging() {
logrus.SetLevel(logrusLevels[viper.GetString("log_level")])
logrus.SetReportCaller(viper.GetBool("log_caller"))
logrus.SetFormatter(&logrus.TextFormatter{
FullTimestamp: true,
TimestampFormat: "2006-01-02T15:04:05.999999999Z07:00",
})
}
func validateOptions() {
if viper.GetString("api_keys_table") == "" {
panic("api_keys_table is required")
}
}
func init() {
initOptions()
initLogging()
validateOptions()
}
func main() {
if len(os.Args) != 2 {
fmt.Println("usage: api-key-generator <user_id>")
os.Exit(1)
}
userName := os.Args[1]
cfg, err := config.LoadDefaultConfig(context.TODO())
if err != nil {
panic(err)
}
dynamodbClient := dynamodb.NewFromConfig(cfg)
// generate random api key
apiKeyBytes := make([]byte, 32)
_, err = rand.Read(apiKeyBytes)
if err != nil {
panic(err)
}
// hash the api key
h := sha256.New()
h.Write(apiKeyBytes)
apiKeyHash := h.Sum(nil)
// check if the api key is already in the table
getResponse, err := dynamodbClient.GetItem(context.TODO(), &dynamodb.GetItemInput{
TableName: aws.String(viper.GetString("api_keys_table")),
Key: map[string]types.AttributeValue{
"api_key_hash": &types.AttributeValueMemberB{
Value: apiKeyHash,
},
},
})
if err != nil {
panic(err)
}
if getResponse.Item == nil {
// api key doesn't exist, add it to dynamodb
_, err := dynamodbClient.PutItem(context.TODO(), &dynamodb.PutItemInput{
TableName: aws.String(viper.GetString("api_keys_table")),
Item: map[string]types.AttributeValue{
"api_key_hash": &types.AttributeValueMemberB{
Value: apiKeyHash,
},
"user_id": &types.AttributeValueMemberS{
Value: userName,
},
},
})
if err != nil {
panic(err)
}
fmt.Println("API key:", base64.StdEncoding.EncodeToString(apiKeyBytes))
} else {
panic("api key already exists, please try again")
}
}