-
Notifications
You must be signed in to change notification settings - Fork 29
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
ffd58eb
commit bf6b6b6
Showing
7 changed files
with
278 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| package server | ||
|
|
||
| import ( | ||
| "crypto/rand" | ||
| "encoding/hex" | ||
| "log" | ||
| "net/http" | ||
| "os" | ||
| ) | ||
|
|
||
| func generateAPIKey() string { | ||
| // if NUCLEARPOND_API_KEY is set, use that | ||
| if os.Getenv("NUCLEARPOND_API_KEY") != "" { | ||
| return os.Getenv("NUCLEARPOND_API_KEY") | ||
| } | ||
| // otherwise, generate a random API key | ||
| b := make([]byte, 16) | ||
| _, err := rand.Read(b) | ||
| if err != nil { | ||
| log.Fatal(err) | ||
| } | ||
| apiKey := hex.EncodeToString(b) | ||
| log.Println("Generated API key:", apiKey) | ||
| os.Setenv("NUCLEARPOND_API_KEY", apiKey) | ||
| return apiKey | ||
| } | ||
|
|
||
| func checkAPIKey(r *http.Request) bool { | ||
| apiKey := generateAPIKey() | ||
| key := r.Header.Get("X-NuclearPond-API-Key") | ||
| if key != apiKey { | ||
| return false | ||
| } | ||
| return true | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,110 @@ | ||
| package server | ||
|
|
||
| import ( | ||
| "encoding/base64" | ||
| "fmt" | ||
| "log" | ||
| "os" | ||
| "strings" | ||
| "time" | ||
|
|
||
| "github.com/DevSecOpsDocs/nuclearpond/pkg/core" | ||
| "github.com/DevSecOpsDocs/nuclearpond/pkg/helpers" | ||
| "github.com/aws/aws-sdk-go/aws" | ||
| "github.com/aws/aws-sdk-go/aws/session" | ||
| "github.com/aws/aws-sdk-go/service/dynamodb" | ||
| ) | ||
|
|
||
| func backgroundScan(scanInput Request, scanId string) { | ||
| targets := helpers.RemoveEmpty(scanInput.Targets) | ||
| batches := helpers.SplitSlice(targets, scanInput.Batches) | ||
| output := scanInput.Output | ||
| threads := scanInput.Threads | ||
| NucleiArgs := base64.StdEncoding.EncodeToString([]byte(scanInput.Args)) | ||
| silent := true | ||
|
|
||
| // Fail if AWS_LAMBDA_FUNCTION_NAME and AWS_REGION are not set | ||
| functionName := os.Getenv("AWS_LAMBDA_FUNCTION_NAME") | ||
| regionName := os.Getenv("AWS_REGION") | ||
| dynamodbTable := os.Getenv("AWS_DYNAMODB_TABLE") | ||
| if functionName == "" || regionName == "" || dynamodbTable == "" { | ||
| log.Fatal("AWS_LAMBDA_FUNCTION_NAME is not set") | ||
| } | ||
|
|
||
| // Convert scanId to a valid DynamoDB key | ||
| requestId := strings.ReplaceAll(scanId, "-", "") | ||
|
|
||
| log.Println("Initiating scan with the id of ", scanId, "with", len(targets), "targets") | ||
| storeScanState(requestId, "running") | ||
| core.ExecuteScans(batches, output, functionName, NucleiArgs, threads, silent) | ||
| storeScanState(requestId, "completed") | ||
| log.Println("Scan", scanId, "completed") | ||
| } | ||
|
|
||
| func storeScanState(requestId string, status string) error { | ||
| log.Println("Stored scan state in Dynamodb", requestId, "as", status) | ||
|
|
||
| sess, err := session.NewSession(&aws.Config{ | ||
| Region: aws.String(os.Getenv("AWS_REGION")), | ||
| }) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| // Create DynamoDB client | ||
| svc := dynamodb.New(sess) | ||
| // Prepare the item to be put into the DynamoDB table | ||
| item := &dynamodb.PutItemInput{ | ||
| TableName: aws.String(os.Getenv("AWS_DYNAMODB_TABLE")), | ||
| Item: map[string]*dynamodb.AttributeValue{ | ||
| "scan_id": { | ||
| S: aws.String(requestId), | ||
| }, | ||
| "status": { | ||
| S: aws.String(status), | ||
| }, | ||
| "timestamp": { | ||
| N: aws.String(fmt.Sprintf("%d", time.Now().Unix())), | ||
| }, | ||
| "ttl": { | ||
| N: aws.String(fmt.Sprintf("%d", time.Now().Add(time.Duration(30*time.Minute)).Unix())), | ||
| }, | ||
| }, | ||
| } | ||
| // Store the item in DynamoDB | ||
| _, err = svc.PutItem(item) | ||
| if err != nil { | ||
| log.Println("Failed to store scan state in Dynamodb:", err) | ||
| return err | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| // function to retrieve the scan state from DynamoDB | ||
| func getScanState(requestId string) (string, error) { | ||
| log.Println("Retrieving scan state from Dynamodb", requestId) | ||
|
|
||
| sess, err := session.NewSession(&aws.Config{ | ||
| Region: aws.String(os.Getenv("AWS_REGION")), | ||
| }) | ||
| if err != nil { | ||
| return "failed", err | ||
| } | ||
| // Create DynamoDB client | ||
| svc := dynamodb.New(sess) | ||
| // Prepare the item to be put into the DynamoDB table | ||
| item := &dynamodb.GetItemInput{ | ||
| TableName: aws.String(os.Getenv("AWS_DYNAMODB_TABLE")), | ||
| Key: map[string]*dynamodb.AttributeValue{ | ||
| "scan_id": { | ||
| S: aws.String(requestId), | ||
| }, | ||
| }, | ||
| } | ||
| // Store the item in DynamoDB | ||
| result, err := svc.GetItem(item) | ||
| if err != nil { | ||
| return "failed", err | ||
| } | ||
| return *result.Item["status"].S, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,94 @@ | ||
| package server | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "fmt" | ||
| "log" | ||
| "net/http" | ||
| "strings" | ||
|
|
||
| "github.com/google/uuid" | ||
| ) | ||
|
|
||
| type Request struct { | ||
| Targets []string `json:"Targets"` | ||
| Batches int `json:"Batches"` | ||
| Threads int `json:"Threads"` | ||
| Args string `json:"Args"` | ||
| Output string `json:"Output"` | ||
| } | ||
|
|
||
| // Index | ||
| func indexHandler(w http.ResponseWriter, r *http.Request) { | ||
| // check if the key matches the generated API key | ||
| if !checkAPIKey(r) { | ||
| http.Error(w, "Invalid API key", http.StatusUnauthorized) | ||
| return | ||
| } | ||
| fmt.Fprintf(w, "Welcome to the API") | ||
| } | ||
|
|
||
| // Health check | ||
| func healthHandler(w http.ResponseWriter, r *http.Request) { | ||
| fmt.Fprintf(w, "ok") | ||
| } | ||
|
|
||
| func scanStatusHandler(w http.ResponseWriter, r *http.Request) { | ||
| // check if the key matches the generated API key | ||
| if !checkAPIKey(r) { | ||
| http.Error(w, "Invalid API key", http.StatusUnauthorized) | ||
| return | ||
| } | ||
|
|
||
| path := r.URL.Path | ||
| parts := strings.Split(path, "/") | ||
| if len(parts) < 3 { | ||
| http.Error(w, "Invalid path", http.StatusBadRequest) | ||
| return | ||
| } | ||
| scanId := parts[2] | ||
|
|
||
| log.Println("Getting scan state for", scanId) | ||
|
|
||
| log.Println("Getting scan state for", scanId) | ||
|
|
||
| state, err := getScanState(scanId) | ||
| if err != nil { | ||
| http.Error(w, "Error getting scan state: "+err.Error(), http.StatusInternalServerError) | ||
| return | ||
| } | ||
| json.NewEncoder(w).Encode(map[string]string{"status": state}) | ||
| } | ||
|
|
||
| // Scan | ||
| func scanHandler(w http.ResponseWriter, r *http.Request) { | ||
| // check if the key matches the generated API key | ||
| if !checkAPIKey(r) { | ||
| http.Error(w, "Invalid API key", http.StatusUnauthorized) | ||
| return | ||
| } | ||
|
|
||
| var req Request | ||
| if err := json.NewDecoder(r.Body).Decode(&req); err != nil { | ||
| http.Error(w, "Error decoding JSON: "+err.Error(), http.StatusBadRequest) | ||
| return | ||
| } | ||
|
|
||
| log.Println("Received request", req) | ||
|
|
||
| scanId := uuid.New().String() | ||
| go backgroundScan(req, scanId) | ||
| json.NewEncoder(w).Encode(map[string]string{"RequestId": scanId}) | ||
| } | ||
|
|
||
| func HandleRequests() { | ||
| // generate API key | ||
| generateAPIKey() | ||
|
|
||
| http.HandleFunc("/", indexHandler) | ||
| http.HandleFunc("/health-check", healthHandler) | ||
| http.HandleFunc("/scan", scanHandler) | ||
| http.HandleFunc("/scan/", scanStatusHandler) | ||
|
|
||
| http.ListenAndServe(":8080", nil) | ||
| } |