From 9b19ca427185cbf0849f5d33cd4d360bbc1d6500 Mon Sep 17 00:00:00 2001
From: LenaHoinkis <LenaHoinkis@users.noreply.github.com>
Date: Wed, 17 Jun 2020 15:47:10 +0200
Subject: [PATCH] Support github enterprise (#43)

* Support github enterprise with NewEnterpriseClient
---
 cmd/github_receiver/main.go      | 13 ++++++++++++-
 cmd/github_receiver/main_test.go |  9 +++++++++
 issues/issues.go                 | 24 ++++++++++++++++++++++++
 issues/issues_test.go            | 15 +++++++++++++++
 4 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/cmd/github_receiver/main.go b/cmd/github_receiver/main.go
index 76a7099..659539b 100644
--- a/cmd/github_receiver/main.go
+++ b/cmd/github_receiver/main.go
@@ -42,6 +42,8 @@ var (
 	authtokenFile   = flagx.FileBytes{}
 	githubOrg       = flag.String("org", "", "The github user or organization name where all repos are found.")
 	githubRepo      = flag.String("repo", "", "The default repository for creating issues when alerts do not include a repo label.")
+	githubBaseURL   = flag.String("enterprise.base-url", "", "The URL of your GitHub Enterprise with API suffix (for example '/api/v3/').")
+	githubUploadURL = flag.String("enterprise.upload-url", "", "The upload URL needs to be set if it differs from the Github Enterprise base URL.")
 	enableAutoClose = flag.Bool("enable-auto-close", false, "Once an alert stops firing, automatically close open issues.")
 	labelOnResolved = flag.String("label-on-resolved", "", "Once an alert stops firing, apply this label.")
 	enableInMemory  = flag.Bool("enable-inmemory", false, "Perform all operations in memory, without using github API.")
@@ -124,9 +126,18 @@ func main() {
 	var client alerts.ReceiverClient
 	if *enableInMemory {
 		client = local.NewClient()
-	} else {
+	} else if *githubBaseURL == "" {
 		client = issues.NewClient(*githubOrg, token, *alertLabel)
+	} else {
+		var err error
+		client, err = issues.NewEnterpriseClient(*githubBaseURL, *githubUploadURL, *githubOrg, token, *alertLabel)
+		if err != nil {
+			fmt.Print(err)
+			osExit(1)
+			return
+		}
 	}
+
 	promSrv := prometheusx.MustServeMetrics()
 	defer promSrv.Close()
 
diff --git a/cmd/github_receiver/main_test.go b/cmd/github_receiver/main_test.go
index 3016bbc..56ec31f 100644
--- a/cmd/github_receiver/main_test.go
+++ b/cmd/github_receiver/main_test.go
@@ -21,6 +21,7 @@ func Test_main(t *testing.T) {
 		authfile     string
 		authtoken    string
 		repo         string
+		baseURL      string
 		titleTmpl    string
 		inmemory     bool
 		expectStatus int
@@ -48,6 +49,13 @@ func Test_main(t *testing.T) {
 			titleTmpl:    "{{ x }}",
 			expectStatus: 1,
 		},
+		{
+			name:         "bad-baseURL",
+			repo:         "fake-repo",
+			authtoken:    "token",
+			baseURL:      "invalidURLEscape%zz",
+			expectStatus: 1,
+		},
 	}
 	flag.CommandLine.SetOutput(ioutil.Discard)
 	for _, tt := range tests {
@@ -60,6 +68,7 @@ func Test_main(t *testing.T) {
 		authtokenFile = []byte(tt.authfile)
 		*githubOrg = "fake-org"
 		*githubRepo = tt.repo
+		*githubBaseURL = tt.baseURL
 		*enableInMemory = tt.inmemory
 		// Guarantee no port conflicts between tests of main.
 		*prometheusx.ListenAddress = ":0"
diff --git a/issues/issues.go b/issues/issues.go
index 4212e54..9bfb219 100644
--- a/issues/issues.go
+++ b/issues/issues.go
@@ -90,6 +90,7 @@ func NewClient(org, authToken, alertLabel string) *Client {
 	tokenSource := oauth2.StaticTokenSource(
 		&oauth2.Token{AccessToken: authToken},
 	)
+
 	client := &Client{
 		GithubClient: github.NewClient(oauth2.NewClient(ctx, tokenSource)),
 		org:          org,
@@ -98,6 +99,29 @@ func NewClient(org, authToken, alertLabel string) *Client {
 	return client
 }
 
+// NewEnterpriseClient creates an Enterprise Client authenticated using the Github authToken.
+// Future operations are only performed on the given github enterprise "org/repo".
+// If uploadURL is empty it will be set to baseURL
+func NewEnterpriseClient(baseURL, uploadURL, org, authToken, alertLabel string) (*Client, error) {
+	ctx := context.Background()
+	tokenSource := oauth2.StaticTokenSource(
+		&oauth2.Token{AccessToken: authToken},
+	)
+
+	if uploadURL == "" {
+		uploadURL = baseURL
+	}
+
+	githubClient, err := github.NewEnterpriseClient(baseURL, uploadURL, oauth2.NewClient(ctx, tokenSource))
+
+	client := &Client{
+		GithubClient: githubClient,
+		org:          org,
+		alertLabel:   alertLabel,
+	}
+	return client, err
+}
+
 // CreateIssue creates a new Github issue. New issues are unassigned. Issues are
 // labeled with with an alert named alertLabel. Labels are created automatically
 // if they do not already exist in a repo.
diff --git a/issues/issues_test.go b/issues/issues_test.go
index 5a19380..5940d6a 100644
--- a/issues/issues_test.go
+++ b/issues/issues_test.go
@@ -420,3 +420,18 @@ func TestClient_rateLimit(t *testing.T) {
 		return
 	}
 }
+
+func TestClient_newEnterpriseClient(t *testing.T) {
+	baseURL := "https://github.example.com/"
+	c, _ := issues.NewEnterpriseClient(baseURL, "", "fake-org", "FAKE-AUTH-TOKEN", "alert")
+
+	if c.GithubClient.BaseURL.String() != baseURL {
+		t.Errorf("client baseURL got %q but want %q", c.GithubClient.BaseURL.String(), baseURL)
+		return
+	}
+
+	if c.GithubClient.UploadURL.String() != baseURL {
+		t.Errorf("client uploadURL got %q but want %q", c.GithubClient.UploadURL.String(), baseURL)
+		return
+	}
+}