From 3cce306928e5c30d2a6d1b51c06296800880a3ed Mon Sep 17 00:00:00 2001
From: Jonathan Pautz <jonathan@datavant.com>
Date: Sat, 29 Jul 2023 14:23:15 -0400
Subject: [PATCH] fix updates for wiz_user and wiz_project, add SAML to
 wiz_user #154 #155

---
 internal/provider/resource_project.go | 82 +++++++++++++++++++++++++--
 internal/provider/resource_user.go    | 40 +++++++++++--
 internal/wiz/structs.go               |  1 +
 3 files changed, 113 insertions(+), 10 deletions(-)

diff --git a/internal/provider/resource_project.go b/internal/provider/resource_project.go
index b10f0c7..f12f933 100644
--- a/internal/provider/resource_project.go
+++ b/internal/provider/resource_project.go
@@ -887,11 +887,83 @@ func resourceWizProjectUpdate(ctx context.Context, d *schema.ResourceData, m int
 
 	// define the graphql query
 	query := `mutation UpdateProject($input: UpdateProjectInput!) {
-	  updateProject(input: $input) {
-	    project {
-	      id
-	    }
-	  }
+		updateProject(input: $input) {
+			project {
+				id
+				name
+				identifiers
+				description
+				businessUnit
+				projectOwners {
+					id
+					name
+					email
+				}
+				securityChampions {
+					id
+					name
+					email
+				}
+				cloudOrganizationLinks {
+					cloudOrganization {
+						id
+					}
+					environment
+					resourceTags {
+						key
+						value
+					}
+					shared
+					resourceGroups
+				}
+				cloudAccountLinks {
+					cloudAccount {
+						id
+					}
+					environment
+					resourceTags {
+						key
+						value
+					}
+					shared
+					resourceGroups
+				}
+				kubernetesClustersLinks {
+					kubernetesCluster {
+						id
+					}
+					environment
+					namespaces
+					shared
+				}
+				repositoryLinks {
+					repository {
+						id
+					}
+				}
+				containerRegistryLinks {
+					containerRegistry {
+						id
+					}
+					environment
+				}
+				ancestorProjects {
+					id
+					name
+				}
+				riskProfile {
+					businessImpact
+					hasAuthentication
+					isInternetFacing
+					hasExposedAPI
+					storesData
+					sensitiveDataTypes
+					regulatoryStandards
+					isCustomerFacing
+					isRegulated
+				}
+			}
+		}
 	}`
 
 	// populate the graphql variables
diff --git a/internal/provider/resource_user.go b/internal/provider/resource_user.go
index 0c42b3b..8ba60a9 100644
--- a/internal/provider/resource_user.go
+++ b/internal/provider/resource_user.go
@@ -34,6 +34,11 @@ func resourceWizUser() *schema.Resource {
 				Description: "The user email address.",
 				Required:    true,
 			},
+			"idpID": {
+				Type:        schema.TypeString,
+				Description: "IDP ID.",
+				Required:    false,
+			},
 			"role": {
 				Type:        schema.TypeString,
 				Description: "Whether the project is archived/inactive",
@@ -86,9 +91,23 @@ func resourceWizUserCreate(ctx context.Context, d *schema.ResourceData, m interf
 
 	// populate the graphql variables
 	vars := &wiz.CreateUserInput{}
+	vars.IdpID = d.Get("idpID").(string)
+
+	// SAML Users require seperate query for create only
+	if ssoEnabled := vars.IdpID != ""; ssoEnabled {
+		query = `mutation CreateSAMLUser($input: CreateSAMLUserInput!) {
+			createSAMLUser(input: $input) {
+				user {
+					id
+				}
+			}
+		}`
+	}
+
 	vars.Name = d.Get("name").(string)
 	vars.Email = d.Get("email").(string)
 	vars.Role = d.Get("role").(string)
+
 	vars.SendEmailInvite = d.Get("send_email_invite").(bool)
 	vars.AssignedProjectIDs = utils.ConvertListToString(d.Get("assigned_project_ids").([]interface{}))
 
@@ -202,11 +221,22 @@ func resourceWizUserUpdate(ctx context.Context, d *schema.ResourceData, m interf
 
 	// define the graphql query
 	query := `mutation UpdateUser($input: UpdateUserInput!) {
-	    updateUser(input: $input) {
-	        user {
-	            id
-	        }
-	    }
+		updateUser(input: $input) {
+			user {
+				id
+				name
+				email
+				isSuspended
+				effectiveRole {
+					id
+					name
+				}
+				effectiveAssignedProjects {
+					id
+					name
+				}
+			}
+		}
 	}`
 
 	// populate the graphql variables
diff --git a/internal/wiz/structs.go b/internal/wiz/structs.go
index d33fd5c..66ef2f1 100644
--- a/internal/wiz/structs.go
+++ b/internal/wiz/structs.go
@@ -1072,6 +1072,7 @@ type CreateUserInput struct {
 	Role               string   `json:"role"`
 	AssignedProjectIDs []string `json:"assignedProjectIds,omitempty"`
 	SendEmailInvite    bool     `json:"sendEmailInvite"`
+	IdpID              string   `json:"idpID,omitempty"`
 }
 
 // CreateUserPayload struct