-
Notifications
You must be signed in to change notification settings - Fork 28
/
main.go
212 lines (182 loc) · 5.41 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
package main
import (
"context"
"flag"
"fmt"
"net/http"
"net/url"
"os"
"os/signal"
"strconv"
"strings"
"syscall"
"time"
"github.com/genuinetools/ghb0t/version"
"github.com/genuinetools/pkg/cli"
"github.com/google/go-github/github"
"github.com/gregjones/httpcache"
"github.com/gregjones/httpcache/diskcache"
"github.com/sirupsen/logrus"
"golang.org/x/oauth2"
)
var (
token string
interval time.Duration
enturl string
lastChecked time.Time
debug bool
)
func main() {
// Create a new cli program.
p := cli.NewProgram()
p.Name = "ghb0t"
p.Description = "A GitHub Bot to automatically delete your fork's branches after a pull request has been merged"
// Set the GitCommit and Version.
p.GitCommit = version.GITCOMMIT
p.Version = version.VERSION
// Setup the global flags.
p.FlagSet = flag.NewFlagSet("global", flag.ExitOnError)
p.FlagSet.StringVar(&token, "token", os.Getenv("GITHUB_TOKEN"), "GitHub API token (or env var GITHUB_TOKEN)")
p.FlagSet.DurationVar(&interval, "interval", 30*time.Second, "check interval (ex. 5ms, 10s, 1m, 3h)")
p.FlagSet.StringVar(&enturl, "url", "", "Connect to a specific GitHub server, provide full API URL (ex. https://github.example.com/api/v3/)")
p.FlagSet.BoolVar(&debug, "d", false, "enable debug logging")
// Set the before function.
p.Before = func(ctx context.Context) error {
// Set the log level.
if debug {
logrus.SetLevel(logrus.DebugLevel)
}
if token == "" {
return fmt.Errorf("GitHub token cannot be empty")
}
return nil
}
// Set the main program action.
p.Action = func(ctx context.Context, args []string) error {
ticker := time.NewTicker(interval)
// On ^C, or SIGTERM handle exit.
signals := make(chan os.Signal, 1)
signal.Notify(signals, os.Interrupt)
signal.Notify(signals, syscall.SIGTERM)
var cancel context.CancelFunc
ctx, cancel = context.WithCancel(ctx)
go func() {
for sig := range signals {
cancel()
ticker.Stop()
logrus.Infof("Received %s, exiting.", sig.String())
os.Exit(0)
}
}()
// Create the http client.
ts := oauth2.StaticTokenSource(
&oauth2.Token{AccessToken: token},
)
// Create the HTTP cache.
cachePath := "/tmp/cache"
if err := os.MkdirAll(cachePath, 0777); err != nil {
logrus.Fatal(err)
}
cache := diskcache.New(cachePath)
tr := httpcache.NewTransport(cache)
c := &http.Client{Transport: tr}
ctx = context.WithValue(ctx, oauth2.HTTPClient, c)
// Create the github client.
tc := oauth2.NewClient(ctx, ts)
client := github.NewClient(tc)
if enturl != "" {
var err error
client.BaseURL, err = url.Parse(enturl + "/api/v3/")
if err != nil {
logrus.Fatal(err)
}
}
// Get the authenticated user, the empty string being passed let's the GitHub
// API know we want ourself.
user, _, err := client.Users.Get(ctx, "")
if err != nil {
logrus.Fatal(err)
}
username := *user.Login
logrus.Infof("Bot started for user %s.", username)
for range ticker.C {
page := 1
perPage := 20
if err := getNotifications(ctx, client, username, page, perPage); err != nil {
logrus.Warn(err)
}
}
return nil
}
// Run our program.
p.Run()
}
// getNotifications iterates over all the notifications received by a user.
func getNotifications(ctx context.Context, client *github.Client, username string, page, perPage int) error {
opt := &github.NotificationListOptions{
All: true,
Since: lastChecked,
ListOptions: github.ListOptions{
Page: page,
PerPage: perPage,
},
}
if lastChecked.IsZero() {
lastChecked = time.Now()
}
notifications, resp, err := client.Activity.ListNotifications(ctx, opt)
if err != nil {
return err
}
for _, notification := range notifications {
// handle event
if err := handleNotification(ctx, client, notification, username); err != nil {
return err
}
}
// Return early if we are on the last page.
if page == resp.LastPage || resp.NextPage == 0 {
return nil
}
page = resp.NextPage
return getNotifications(ctx, client, username, page, perPage)
}
func handleNotification(ctx context.Context, client *github.Client, notification *github.Notification, username string) error {
// Check if the type is a pull request.
if *notification.Subject.Type == "PullRequest" {
// Let's get some information about the pull request.
parts := strings.Split(*notification.Subject.URL, "/")
last := parts[len(parts)-1]
id, err := strconv.Atoi(last)
if err != nil {
return err
}
pr, _, err := client.PullRequests.Get(ctx, *notification.Repository.Owner.Login, *notification.Repository.Name, id)
if err != nil {
return err
}
if *pr.State == "closed" && *pr.Merged {
// If the PR was made from a repository owned by the current user,
// let's delete it.
branch := *pr.Head.Ref
if pr.Head.Repo == nil {
return nil
}
if pr.Head.Repo.Owner == nil {
return nil
}
owner := *pr.Head.Repo.Owner.Login
defaultBranch := *notification.Repository.DefaultBranch
// Never delete the default branch or a branch we do not own.
if owner == username && branch != defaultBranch {
_, err := client.Git.DeleteRef(ctx, username, *pr.Head.Repo.Name, strings.Replace("heads/"+*pr.Head.Ref, "#", "%23", -1))
// 422 is the error code for when the branch does not exist.
if err != nil && !strings.Contains(err.Error(), " 422 ") {
return err
}
logrus.Infof("Branch %s on %s/%s no longer exists.", branch, owner, *pr.Head.Repo.Name)
}
}
}
return nil
}