From 27e855ddbde2feb9e4d07b21c0f613882cb46dfb Mon Sep 17 00:00:00 2001
From: Kohsuke Kawaguchi <kk@kohsuke.org>
Date: Thu, 2 Jun 2016 23:40:37 -0700
Subject: [PATCH] Issue 262: added support for branch protection API

---
 .../org/kohsuke/github/BranchProtection.java  | 18 ++++++++
 .../org/kohsuke/github/EnforcementLevel.java  | 14 +++++++
 .../java/org/kohsuke/github/GHBranch.java     | 42 +++++++++++++++++++
 .../java/org/kohsuke/github/GHRepository.java |  4 ++
 .../github/GHContentIntegrationTest.java      |  7 ++++
 5 files changed, 85 insertions(+)
 create mode 100644 src/main/java/org/kohsuke/github/BranchProtection.java
 create mode 100644 src/main/java/org/kohsuke/github/EnforcementLevel.java

diff --git a/src/main/java/org/kohsuke/github/BranchProtection.java b/src/main/java/org/kohsuke/github/BranchProtection.java
new file mode 100644
index 0000000000..edeaa4d86c
--- /dev/null
+++ b/src/main/java/org/kohsuke/github/BranchProtection.java
@@ -0,0 +1,18 @@
+package org.kohsuke.github;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * @author Kohsuke Kawaguchi
+ * @see GHBranch#disableProtection()
+ */
+class BranchProtection {
+    boolean enabled;
+    RequiredStatusChecks requiredStatusChecks;
+
+    static class RequiredStatusChecks {
+        EnforcementLevel enforcement_level;
+        List<String> contexts = new ArrayList<String>();
+    }
+}
diff --git a/src/main/java/org/kohsuke/github/EnforcementLevel.java b/src/main/java/org/kohsuke/github/EnforcementLevel.java
new file mode 100644
index 0000000000..1b29195299
--- /dev/null
+++ b/src/main/java/org/kohsuke/github/EnforcementLevel.java
@@ -0,0 +1,14 @@
+package org.kohsuke.github;
+
+import java.util.Locale;
+
+/**
+ * @author Kohsuke Kawaguchi
+ */
+public enum EnforcementLevel {
+    OFF, NON_ADMINS, EVERYONE;
+
+    public String toString() {
+        return name().toLowerCase(Locale.ENGLISH);
+    }
+}
diff --git a/src/main/java/org/kohsuke/github/GHBranch.java b/src/main/java/org/kohsuke/github/GHBranch.java
index 272c643f47..5f41dd8e91 100644
--- a/src/main/java/org/kohsuke/github/GHBranch.java
+++ b/src/main/java/org/kohsuke/github/GHBranch.java
@@ -1,6 +1,11 @@
 package org.kohsuke.github;
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+import org.kohsuke.github.BranchProtection.RequiredStatusChecks;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collection;
 
 /**
  * A branch in a repository.
@@ -44,6 +49,43 @@ public String getName() {
     public String getSHA1() {
         return commit.sha;
     }
+
+    /**
+     * Disables branch protection and allows anyone with push access to push changes.
+     */
+    public void disableProtection() throws IOException {
+        BranchProtection bp = new BranchProtection();
+        bp.enabled = false;
+        setProtection(bp);
+    }
+
+    /**
+     * Enables branch protection to control what commit statuses are required to push.
+     *
+     * @see GHCommitStatus#getContext()
+     */
+    public void enableProtection(EnforcementLevel level, Collection<String> contexts) throws IOException {
+        BranchProtection bp = new BranchProtection();
+        bp.enabled = true;
+        bp.requiredStatusChecks = new RequiredStatusChecks();
+        bp.requiredStatusChecks.enforcement_level = level;
+        bp.requiredStatusChecks.contexts.addAll(contexts);
+        setProtection(bp);
+    }
+
+    public void enableProtection(EnforcementLevel level, String... contexts) throws IOException {
+        enableProtection(level, Arrays.asList(contexts));
+    }
+
+    private void setProtection(BranchProtection bp) throws IOException {
+        new Requester(root).method("PATCH")
+                .withHeader("Accept","application/vnd.github.loki-preview+json")
+                ._with("protection",bp).to(getApiRoute());
+    }
+
+    String getApiRoute() {
+        return owner.getApiTailUrl("/branches/"+name);
+    }
     
     @Override
     public String toString() {
diff --git a/src/main/java/org/kohsuke/github/GHRepository.java b/src/main/java/org/kohsuke/github/GHRepository.java
index 7afd59e4dc..dab8166cc3 100644
--- a/src/main/java/org/kohsuke/github/GHRepository.java
+++ b/src/main/java/org/kohsuke/github/GHRepository.java
@@ -1103,6 +1103,10 @@ public Map<String,GHBranch> getBranches() throws IOException {
         return r;
     }
 
+    public GHBranch getBranch(String name) throws IOException {
+        return root.retrieve().to(getApiTailUrl("branches/"+name),GHBranch.class).wrap(this);
+    }
+
     /**
      * @deprecated
      *      Use {@link #listMilestones(GHIssueState)}
diff --git a/src/test/java/org/kohsuke/github/GHContentIntegrationTest.java b/src/test/java/org/kohsuke/github/GHContentIntegrationTest.java
index 8e3873708b..c642d662e7 100644
--- a/src/test/java/org/kohsuke/github/GHContentIntegrationTest.java
+++ b/src/test/java/org/kohsuke/github/GHContentIntegrationTest.java
@@ -20,6 +20,13 @@ public void setUp() throws Exception {
         repo = gitHub.getRepository("github-api-test-org/GHContentIntegrationTest").fork();
     }
 
+    @Test
+    public void testBranchProtection() throws Exception {
+        GHBranch b = repo.getBranch("master");
+        b.enableProtection(EnforcementLevel.NON_ADMINS, "foo/bar");
+        b.disableProtection();
+    }
+
     @Test
     public void testGetFileContent() throws Exception {
         GHContent content = repo.getFileContent("ghcontent-ro/a-file-with-content");