-
Notifications
You must be signed in to change notification settings - Fork 35
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #307 from projectdiscovery/pdcp-cred-helpers
Add pdcp credential helpers
- Loading branch information
Showing
5 changed files
with
359 additions
and
11 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| // pdcp contains projectdiscovery cloud platform related features | ||
| // like result upload , dashboard etc. | ||
| package pdcp | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "os" | ||
| "strings" | ||
|
|
||
| "github.com/projectdiscovery/gologger" | ||
| "github.com/projectdiscovery/utils/env" | ||
| "golang.org/x/term" | ||
| ) | ||
|
|
||
| var ( | ||
| DashBoardURL = "https://cloud.projectdiscovery.io" | ||
| DefaultApiServer = "https://api.projectdiscovery.io" | ||
| ) | ||
|
|
||
| // CheckNValidateCredentials checks if credentials exist on filesystem | ||
| // if not waits for user to enter credentials and validates them | ||
| // and saves them to filesystem | ||
| // when validate is true any existing credentials are validated | ||
| // Note: this is meant to be used in cli only (interactive mode) | ||
| func CheckNValidateCredentials(toolName string) { | ||
| h := &PDCPCredHandler{} | ||
| creds, err := h.GetCreds() | ||
| if err == nil { | ||
| // validate by fetching user profile | ||
| gotCreds, err := h.ValidateAPIKey(creds.APIKey, creds.Server, toolName) | ||
| if err == nil { | ||
| gologger.Info().Msgf("You are logged in as (@%v)", gotCreds.Username) | ||
| os.Exit(0) | ||
| } | ||
| gologger.Error().Msgf("Invalid API key found in file, please recheck or recreate your API key and retry.") | ||
| } | ||
| if err != nil && err != ErrNoCreds { | ||
| // this is unexpected error log it | ||
| gologger.Error().Msgf("Could not read credentials from file: %s\n", err) | ||
| } | ||
|
|
||
| // if we are here, we need to get credentials from user | ||
| gologger.Info().Msgf("Get your free api key by signing up at %v", DashBoardURL) | ||
| fmt.Printf("[*] Enter PDCP API Key (exit to abort): ") | ||
| bin, err := term.ReadPassword(int(os.Stdin.Fd())) | ||
| if err != nil { | ||
| gologger.Fatal().Msgf("Could not read input from terminal: %s\n", err) | ||
| } | ||
| apiKey := string(bin) | ||
| if strings.EqualFold(apiKey, "exit") { | ||
| os.Exit(0) | ||
| } | ||
| fmt.Println() | ||
| // if env variable is set use that for validating api key | ||
| apiServer := env.GetEnvOrDefault(apiServerEnv, DefaultApiServer) | ||
| // validate by fetching user profile | ||
| validatedCreds, err := h.ValidateAPIKey(apiKey, apiServer, toolName) | ||
| if err == nil { | ||
| gologger.Info().Msgf("Successfully logged in as (@%v)", validatedCreds.Username) | ||
| if saveErr := h.SaveCreds(validatedCreds); saveErr != nil { | ||
| gologger.Warning().Msgf("Could not save credentials to file: %s\n", saveErr) | ||
| } | ||
| os.Exit(0) | ||
| } | ||
| gologger.Error().Msgf("Invalid API key '%v' got error: %v", maskKey(apiKey), err) | ||
| gologger.Fatal().Msgf("please recheck or recreate your API key and retry") | ||
| } | ||
|
|
||
| func maskKey(key string) string { | ||
| if len(key) < 6 { | ||
| // this is invalid key | ||
| return key | ||
| } | ||
| return fmt.Sprintf("%v%v", key[:3], strings.Repeat("*", len(key)-3)) | ||
| } |
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,139 @@ | ||
| package pdcp | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "fmt" | ||
| "io" | ||
| "os" | ||
| "path/filepath" | ||
|
|
||
| "github.com/projectdiscovery/retryablehttp-go" | ||
| "github.com/projectdiscovery/utils/env" | ||
| fileutil "github.com/projectdiscovery/utils/file" | ||
| folderutil "github.com/projectdiscovery/utils/folder" | ||
| urlutil "github.com/projectdiscovery/utils/url" | ||
| "gopkg.in/yaml.v3" | ||
| ) | ||
|
|
||
| var ( | ||
| PDCPDir = filepath.Join(folderutil.HomeDirOrDefault(""), ".pdcp") | ||
| PDCPCredFile = filepath.Join(PDCPDir, "credentials.yaml") | ||
| ErrNoCreds = fmt.Errorf("no credentials found in %s", PDCPDir) | ||
| ) | ||
|
|
||
| const ( | ||
| userProfileURL = "https://%s/v1/user?utm_source=%s" | ||
| apiKeyEnv = "PDCP_API_KEY" | ||
| apiServerEnv = "PDCP_API_SERVER" | ||
| ApiKeyHeaderName = "X-Api-Key" | ||
| dashBoardEnv = "PDCP_DASHBOARD_URL" | ||
| ) | ||
|
|
||
| type PDCPCredentials struct { | ||
| Username string `yaml:"username"` | ||
| APIKey string `yaml:"api-key"` | ||
| Server string `yaml:"server"` | ||
| } | ||
|
|
||
| type PDCPUserProfileResponse struct { | ||
| UserName string `json:"name"` | ||
| // there are more fields but we don't need them | ||
| /// below fields are added later on and not part of the response | ||
| } | ||
|
|
||
| // PDCPCredHandler is interface for adding / retrieving pdcp credentials | ||
| // from file system | ||
| type PDCPCredHandler struct{} | ||
|
|
||
| // GetCreds retrieves the credentials from the file system or environment variables | ||
| func (p *PDCPCredHandler) GetCreds() (*PDCPCredentials, error) { | ||
| credsFromEnv := p.getCredsFromEnv() | ||
| if credsFromEnv != nil { | ||
| return credsFromEnv, nil | ||
| } | ||
| if !fileutil.FolderExists(PDCPDir) || !fileutil.FileExists(PDCPCredFile) { | ||
| return nil, ErrNoCreds | ||
| } | ||
| bin, err := os.Open(PDCPCredFile) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| // for future use-cases | ||
| var creds []PDCPCredentials | ||
| err = yaml.NewDecoder(bin).Decode(&creds) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| if len(creds) == 0 { | ||
| return nil, ErrNoCreds | ||
| } | ||
| return &creds[0], nil | ||
| } | ||
|
|
||
| // getCredsFromEnv retrieves the credentials from the environment | ||
| // if not or incomplete credentials are found it return nil | ||
| func (p *PDCPCredHandler) getCredsFromEnv() *PDCPCredentials { | ||
| apiKey := env.GetEnvOrDefault(apiKeyEnv, "") | ||
| apiServer := env.GetEnvOrDefault(apiServerEnv, "") | ||
| if apiKey == "" || apiServer == "" { | ||
| return nil | ||
| } | ||
| return &PDCPCredentials{APIKey: apiKey, Server: apiServer} | ||
| } | ||
|
|
||
| // SaveCreds saves the credentials to the file system | ||
| func (p *PDCPCredHandler) SaveCreds(resp *PDCPCredentials) error { | ||
| if resp == nil { | ||
| return fmt.Errorf("invalid response") | ||
| } | ||
| if !fileutil.FolderExists(PDCPDir) { | ||
| _ = fileutil.CreateFolder(PDCPDir) | ||
| } | ||
| bin, err := yaml.Marshal([]*PDCPCredentials{resp}) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| return os.WriteFile(PDCPCredFile, bin, 0600) | ||
| } | ||
|
|
||
| // ValidateAPIKey validates the api key and retrieves associated user metadata like username | ||
| // from given api server/host | ||
| func (p *PDCPCredHandler) ValidateAPIKey(key string, host string, toolName string) (*PDCPCredentials, error) { | ||
| // get address from url | ||
| urlx, err := urlutil.Parse(host) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| req, err := retryablehttp.NewRequest("GET", fmt.Sprintf(userProfileURL, urlx.Host, toolName), nil) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| req.Header.Set(ApiKeyHeaderName, key) | ||
| resp, err := retryablehttp.DefaultHTTPClient.Do(req) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| if resp.StatusCode != 200 { | ||
| _, _ = io.Copy(io.Discard, resp.Body) | ||
| _ = resp.Body.Close() | ||
| return nil, fmt.Errorf("invalid status code: %d", resp.StatusCode) | ||
| } | ||
| defer resp.Body.Close() | ||
| bin, err := io.ReadAll(resp.Body) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| var profile PDCPUserProfileResponse | ||
| err = json.Unmarshal(bin, &profile) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| if profile.UserName == "" { | ||
| return nil, fmt.Errorf("invalid response from server got %v", string(bin)) | ||
| } | ||
| return &PDCPCredentials{Username: profile.UserName, APIKey: key, Server: host}, nil | ||
| } | ||
|
|
||
| func init() { | ||
| DashBoardURL = env.GetEnvOrDefault("PDCP_DASHBOARD_URL", DashBoardURL) | ||
| } |
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,33 @@ | ||
| package pdcp | ||
|
|
||
| import ( | ||
| "os" | ||
| "path/filepath" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| var exampleCred = ` | ||
| - username: test | ||
| api-key: testpassword | ||
| server: https://scanme.sh | ||
| ` | ||
|
|
||
| func TestLoadCreds(t *testing.T) { | ||
| // temporarily change PDCP file location for testing | ||
| f, err := os.CreateTemp("", "creds-test-*") | ||
| require.Nil(t, err) | ||
| _, _ = f.WriteString(strings.TrimSpace(exampleCred)) | ||
| defer os.Remove(f.Name()) | ||
| PDCPCredFile = f.Name() | ||
| PDCPDir = filepath.Dir(f.Name()) | ||
| h := &PDCPCredHandler{} | ||
| value, err := h.GetCreds() | ||
| require.Nil(t, err) | ||
| require.NotNil(t, value) | ||
| require.Equal(t, "test", value.Username) | ||
| require.Equal(t, "testpassword", value.APIKey) | ||
| require.Equal(t, "https://scanme.sh", value.Server) | ||
| } |
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
Oops, something went wrong.