github.com/munnerz/test-infra@v0.0.0-20190108210205-ce3d181dc989/robots/pr-labeler/main.go (about)

     1  /*
     2  Copyright 2017 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // PR labeler provides a way to add a missing ok-to-test label on trusted PRs.
    18  //
    19  // The --token-path determines who interacts with github.
    20  // By default PR labeler runs in dry mode, add --confirm to make it leave comments.
    21  package main
    22  
    23  import (
    24  	"flag"
    25  	"log"
    26  	"net/url"
    27  
    28  	"k8s.io/apimachinery/pkg/util/sets"
    29  	"k8s.io/test-infra/prow/config/secret"
    30  	"k8s.io/test-infra/prow/flagutil"
    31  	"k8s.io/test-infra/prow/github"
    32  )
    33  
    34  const (
    35  	needsOkToTest = "needs-ok-to-test"
    36  	okToTest      = "ok-to-test"
    37  )
    38  
    39  type client interface {
    40  	AddLabel(org, repo string, number int, label string) error
    41  	GetIssueLabels(org, repo string, number int) ([]github.Label, error)
    42  	GetPullRequests(org, repo string) ([]github.PullRequest, error)
    43  	ListCollaborators(org, repo string) ([]github.User, error)
    44  	ListOrgMembers(org, role string) ([]github.TeamMember, error)
    45  }
    46  
    47  type options struct {
    48  	confirm            bool
    49  	endpoint           flagutil.Strings
    50  	org                string
    51  	repo               string
    52  	tokenPath          string
    53  	trustCollaborators bool
    54  }
    55  
    56  func flagOptions() options {
    57  	o := options{
    58  		endpoint: flagutil.NewStrings("https://api.github.com"),
    59  	}
    60  	flag.BoolVar(&o.confirm, "confirm", false, "Mutate github if set")
    61  	flag.StringVar(&o.org, "org", "", "github org")
    62  	flag.StringVar(&o.repo, "repo", "", "github repo")
    63  	flag.StringVar(&o.tokenPath, "token-path", "", "Path to github token")
    64  	flag.BoolVar(&o.trustCollaborators, "trust-collaborators", false, "Also trust PRs from collaborators")
    65  	flag.Parse()
    66  	return o
    67  }
    68  
    69  func main() {
    70  	log.SetFlags(log.LstdFlags | log.Lshortfile)
    71  	o := flagOptions()
    72  
    73  	if o.org == "" {
    74  		log.Fatal("empty --org")
    75  	}
    76  	if o.repo == "" {
    77  		log.Fatal("empty --repo")
    78  	}
    79  	if o.tokenPath == "" {
    80  		log.Fatal("empty --token-path")
    81  	}
    82  
    83  	secretAgent := &secret.Agent{}
    84  	if err := secretAgent.Start([]string{o.tokenPath}); err != nil {
    85  		log.Fatalf("Error starting secrets agent: %v", err)
    86  	}
    87  
    88  	var err error
    89  	for _, ep := range o.endpoint.Strings() {
    90  		_, err = url.ParseRequestURI(ep)
    91  		if err != nil {
    92  			log.Fatalf("Invalid --endpoint URL %q: %v.", ep, err)
    93  		}
    94  	}
    95  
    96  	var c client
    97  	if o.confirm {
    98  		c = github.NewClient(secretAgent.GetTokenGenerator(o.tokenPath), o.endpoint.Strings()...)
    99  	} else {
   100  		c = github.NewDryRunClient(secretAgent.GetTokenGenerator(o.tokenPath), o.endpoint.Strings()...)
   101  	}
   102  
   103  	// get all open PRs
   104  	prs, err := c.GetPullRequests(o.org, o.repo)
   105  	if err != nil {
   106  		log.Fatal(err)
   107  	}
   108  
   109  	// get the list of authors to skip once and use a set for lookups
   110  	skipAuthors := sets.NewString()
   111  
   112  	// skip org members
   113  	members, err := c.ListOrgMembers(o.org, "all")
   114  	if err != nil {
   115  		log.Fatal(err)
   116  	}
   117  	for _, member := range members {
   118  		skipAuthors.Insert(member.Login)
   119  	}
   120  
   121  	// eventually also skip collaborators
   122  	if o.trustCollaborators {
   123  		collaborators, err := c.ListCollaborators(o.org, o.repo)
   124  		if err != nil {
   125  			log.Fatal(err)
   126  		}
   127  		for _, collaborator := range collaborators {
   128  			skipAuthors.Insert(collaborator.Login)
   129  		}
   130  	}
   131  
   132  	for _, pr := range prs {
   133  		// skip PRs from these authors
   134  		if skipAuthors.Has(pr.User.Login) {
   135  			continue
   136  		}
   137  		// skip PRs with *ok-to-test labels
   138  		labels, err := c.GetIssueLabels(o.org, o.repo, pr.Number)
   139  		if err != nil {
   140  			log.Fatal(err)
   141  		}
   142  		if github.HasLabel(okToTest, labels) || github.HasLabel(needsOkToTest, labels) {
   143  			continue
   144  		}
   145  		// only add ok-to-test with --confirm
   146  		if !o.confirm {
   147  			log.Println("Use --confirm to add", okToTest, "to", pr.HTMLURL)
   148  			continue
   149  		}
   150  		if err := c.AddLabel(o.org, o.repo, pr.Number, okToTest); err != nil {
   151  			log.Fatal(err)
   152  		}
   153  	}
   154  }