Skip to content

Commit

Permalink
refactor(internal/ai): Update GeneratePolicy to include serviceName p…
Browse files Browse the repository at this point in the history
…arameter
tolgaOzen committed Dec 27, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 6fcd260 commit 9f5ed7f
Showing 7 changed files with 45 additions and 20 deletions.
20 changes: 15 additions & 5 deletions internal/ai/policy_generate.go
Original file line number Diff line number Diff line change
@@ -211,21 +211,31 @@ func toStringSlice(v interface{}) []string {
}
}

func GeneratePolicy(apiKey, prompt string, resourceArn *string) (IAMPolicy, error) {
func GeneratePolicy(apiKey, prompt string, serviceName, resourceArn *string) (IAMPolicy, error) {
url := "https://api.openai.com/v1/chat/completions"

// Include resourceArn information in the user prompt.
resourceArnInfo := ""
// Build detailed information for the service and resource.
serviceAndResourceDetails := ""
if resourceArn != nil {
resourceArnInfo = fmt.Sprintf("\nThe resource ARN is: %s", *resourceArn)
if serviceName != nil {
serviceAndResourceDetails = fmt.Sprintf("The service name is: %s\nThe resource ARN is: %s", *serviceName, *resourceArn)
} else {
serviceAndResourceDetails = fmt.Sprintf("The service name is: all services\nThe resource ARN is: %s", *resourceArn)
}
} else {
if serviceName != nil {
serviceAndResourceDetails = fmt.Sprintf("The service name is: %s\nNo specific resource ARN provided.", *serviceName)
} else {
serviceAndResourceDetails = "The service name is: all services\nNo specific resource ARN provided."
}
}

payload := map[string]interface{}{
"model": "gpt-4o",
"temperature": 0.1,
"messages": []map[string]string{
{"role": "system", "content": "You are an assistant that produces IAM policies as JSON."},
{"role": "user", "content": fmt.Sprintf("%s%s", prompt, resourceArnInfo)},
{"role": "user", "content": fmt.Sprintf("%s%s", prompt, serviceAndResourceDetails)},
},
"response_format": map[string]interface{}{
"type": "json_schema",
4 changes: 3 additions & 1 deletion pkg/aws/groups/controller.go
Original file line number Diff line number Diff line change
@@ -96,7 +96,9 @@ type ResourceLoadedMsg struct{ List []list.Item }
// LoadResources loads resources.
func (c *Controller) LoadResources() tea.Cmd {
return func() tea.Msg {
var items []list.Item
items := []list.Item{
models.Resource{Name: "All Resources", Arn: "*"},
}

resources, err := c.api.ListResources(c.State.GetService().Name)
if err != nil {
12 changes: 7 additions & 5 deletions pkg/aws/groups/create_policy.go
Original file line number Diff line number Diff line change
@@ -80,18 +80,20 @@ func (m CreatePolicy) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return Switch(m.controller.Next(), 0, 0)
} else {
var resourceArn *string = nil
if m.controller.State.GetService() != nil {
if m.controller.State.GetResource() != nil {
resourceArn = &m.controller.State.GetResource().Arn
}

var serviceName *string = nil
if m.controller.State.GetService() != nil {
serviceName = &m.controller.State.GetService().Name
}

if m.message == nil {
m.err = errors.New("Please provide a message")
}

policy, err := ai.GeneratePolicy(m.controller.openAiApiKey, *m.message, resourceArn)
if err != nil {
m.err = err
}
policy, err := ai.GeneratePolicy(m.controller.openAiApiKey, *m.message, serviceName, resourceArn)

policyJson, err := json.MarshalIndent(policy, "", "\t")
if err != nil {
4 changes: 3 additions & 1 deletion pkg/aws/roles/controller.go
Original file line number Diff line number Diff line change
@@ -97,7 +97,9 @@ type ResourceLoadedMsg struct{ List []list.Item }
// LoadResources loads resources.
func (c *Controller) LoadResources() tea.Cmd {
return func() tea.Msg {
var items []list.Item
items := []list.Item{
models.Resource{Name: "All Resources", Arn: "*"},
}

resources, err := c.api.ListResources(c.State.GetService().Name)
if err != nil {
9 changes: 7 additions & 2 deletions pkg/aws/roles/create_policy.go
Original file line number Diff line number Diff line change
@@ -80,15 +80,20 @@ func (m CreatePolicy) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return Switch(m.controller.Next(), 0, 0)
} else {
var resourceArn *string = nil
if m.controller.State.GetService() != nil {
if m.controller.State.GetResource() != nil {
resourceArn = &m.controller.State.GetResource().Arn
}

var serviceName *string = nil
if m.controller.State.GetService() != nil {
serviceName = &m.controller.State.GetService().Name
}

if m.message == nil {
m.err = errors.New("Please provide a message")
}

policy, err := ai.GeneratePolicy(m.controller.openAiApiKey, *m.message, resourceArn)
policy, err := ai.GeneratePolicy(m.controller.openAiApiKey, *m.message, serviceName, resourceArn)
if err != nil {
m.err = err
}
4 changes: 3 additions & 1 deletion pkg/aws/users/controller.go
Original file line number Diff line number Diff line change
@@ -140,7 +140,9 @@ type ResourceLoadedMsg struct{ List []list.Item }
// LoadResources loads resources.
func (c *Controller) LoadResources() tea.Cmd {
return func() tea.Msg {
var items []list.Item
items := []list.Item{
models.Resource{Name: "All Resources", Arn: "*"},
}

resources, err := c.api.ListResources(c.State.GetService().Name)
if err != nil {
12 changes: 7 additions & 5 deletions pkg/aws/users/create_policy.go
Original file line number Diff line number Diff line change
@@ -79,18 +79,20 @@ func (m CreatePolicy) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return Switch(m.controller.Next(), 0, 0)
} else {
var resourceArn *string = nil
if m.controller.State.GetService() != nil {
if m.controller.State.GetResource() != nil {
resourceArn = &m.controller.State.GetResource().Arn
}

var serviceName *string = nil
if m.controller.State.GetService() != nil {
serviceName = &m.controller.State.GetService().Name
}

if m.message == nil {
m.err = errors.New("Please provide a message")
}

policy, err := ai.GeneratePolicy(m.controller.openAiApiKey, *m.message, resourceArn)
if err != nil {
m.err = err
}
policy, err := ai.GeneratePolicy(m.controller.openAiApiKey, *m.message, serviceName, resourceArn)

policyJson, err := json.MarshalIndent(policy, "", "\t")
if err != nil {

0 comments on commit 9f5ed7f

Please sign in to comment.