Adding regex support in the IP whitelist

src/main/java/com/asquera/elasticsearch/plugins/http/auth/InetAddressWhitelist.java

@@ -1,14 +1,16 @@
 package com.asquera.elasticsearch.plugins.http.auth;
-import org.elasticsearch.common.logging.Loggers;
-
+import java.net.Inet4Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
 import java.util.ArrayList;
-import java.util.List;
+import java.util.Arrays;
 import java.util.HashSet;
-import java.util.Set;
 import java.util.Iterator;
-import java.util.Arrays;
-import java.net.InetAddress;
-import java.net.UnknownHostException;
+import java.util.List;
+import java.util.Set;
+import java.util.regex.Pattern;
+
+import org.elasticsearch.common.logging.Loggers;
 
 /**
  *
@@ -22,13 +24,13 @@
  */
 
 public class InetAddressWhitelist {
-  private Set<InetAddress> whitelist;
+  private Set<Object> whitelist;
   /**
    *
    *
    * @param whitelist
    */
-  public InetAddressWhitelist(Set<InetAddress> whitelist) {
+  public InetAddressWhitelist(Set<Object> whitelist) {
     this.whitelist = whitelist;
   }
 
@@ -51,7 +53,21 @@ public InetAddressWhitelist(String[] sWhitelist) {
    * @return if the ip is included in the whitelist
    */
   public Boolean contains(InetAddress candidate) {
-    return this.whitelist.contains(candidate);
+    if (this.whitelist.contains(candidate)){
+    	return true;
+    }
+
+    //We also need to itterate through each of the patterns to make sure it doesn't match there
+    for (Object obj : whitelist){
+   	 if (obj.getClass() == Pattern.class){
+   		 Pattern pattern = (Pattern)obj;
+   		 if (pattern.matcher(candidate.getHostAddress()).matches()){
+   			 return true;
+   		 }
+   	 }
+    }
+
+    return false;
   }
 
   /**
@@ -65,18 +81,38 @@ public Boolean contains(InetAddress candidate) {
    * whitelist ips
    */
   public Boolean contains(String candidate) {
-    return getStringWhitelist().contains(candidate);
+     if (getStringWhitelist().contains(candidate)){
+    	 return true;
+     }
+
+     //We also need to itterate through each of the patterns to make sure it doesn't match there
+     for (Object obj : whitelist){
+    	 if (obj.getClass() == Pattern.class){
+    		 Pattern pattern = (Pattern)obj;
+    		 if (pattern.matcher(candidate).matches()){
+    			 return true;
+    		 }
+    	 }
+     }
+
+     return false;
   }
 
   /**
    * @return set of the string representations of the whitelist
    */
   Set<String> getStringWhitelist() {
-    Iterator<InetAddress> iterator = this.whitelist.iterator();
+    Iterator<Object> iterator = this.whitelist.iterator();
     Set<String> set = new HashSet<String>();
     while (iterator.hasNext()) {
-      InetAddress next = iterator.next();
-      set.add(next.getHostAddress());
+      Object next = iterator.next();
+      if (next.getClass() == Pattern.class){
+    	  set.add(next.toString());
+      } 
+      else{
+    	  InetAddress add = (InetAddress)next;
+    	  set.add(add.getHostAddress());
+      }
     }
     return set;
   }
@@ -89,20 +125,26 @@ Set<String> getStringWhitelist() {
    * @return a list of {@link InetAddress} objects
    *
    */
-  static Set<InetAddress> toInetAddress(List<String> ips) {
-    List<InetAddress> listIps = new ArrayList<InetAddress>();
+  static Set<Object> toInetAddress(List<String> ips) {
+    List<Object> listIps = new ArrayList<Object>();
     Iterator<String> iterator = ips.iterator();
     while (iterator.hasNext()) {
       String next = iterator.next();
-      try {
-        listIps.add(InetAddress.getByName(next));
-      } catch (UnknownHostException e) {
-        String template = "an ip set in the whitelist settings raised an " +
-          "UnknownHostException: {}, dropping it";
-        Loggers.getLogger(InetAddressWhitelist.class).info(template, e.getMessage());
+      if (next != null && next.startsWith("~")){
+    	 Pattern pattern = Pattern.compile(next.substring(1));
+    	 listIps.add(pattern);
+      }
+      else {
+	      try {
+	        listIps.add(InetAddress.getByName(next));
+	      } catch (UnknownHostException e) {
+	        String template = "an ip set in the whitelist settings raised an " +
+	          "UnknownHostException: {}, dropping it";
+	        Loggers.getLogger(InetAddressWhitelist.class).info(template, e.getMessage());
+	      }
       }
     }
-    return new HashSet<InetAddress>(listIps);
+    return new HashSet<Object>(listIps);
   }
 
   /**
@@ -113,4 +155,4 @@ public String toString() {
     return whitelist.toString();
   }
 
-}
+}

src/test/java/com/asquera/elasticsearch/plugins/http/auth/InetAddressWhitelistTest.java

@@ -12,6 +12,7 @@ public class InetAddressWhitelistTest {
   static final String localhost = "localhost";
   static final String containedIp = "1.1.1.1";
   static String notContainedIp = "2.2.2.2";
+  static String containedRegex = "~1.1.*";
   private InetAddressWhitelist whitelist(String ip) {
     String[] w = { ip };
     return new InetAddressWhitelist(w);
@@ -21,10 +22,12 @@ private InetAddressWhitelist whitelist(String ip) {
   public void testInnetLocalhost() throws UnknownHostException {
     assertTrue(whitelist(localhost).contains(InetAddress.getByName(localhost)));
   }
+
   @Test
   public void testInnetNullDefaultsToLocalhost() throws UnknownHostException {
     assertTrue(whitelist(null).contains(InetAddress.getByName(localhost)));
   }
+
   @Test
   public void testStringLocalhostNotMatched() throws UnknownHostException {
     // the ip that "localhost" resolves to its matched ip and not the string
@@ -46,10 +49,21 @@ public void testEmptyWhitelist() throws UnknownHostException {
   public void testNotContained() throws UnknownHostException {
     assertFalse(whitelist(containedIp).contains(notContainedIp));
   }
-
+  
   @Test
   public void invalidIpIsDropped() throws UnknownHostException {
     String invalidIp = "555.555.555.555";
     assertFalse(whitelist(invalidIp).contains(invalidIp));
   }
+
+  @Test 
+  public void testRegexContained() throws UnknownHostException {
+	  assertTrue(whitelist(containedRegex).contains(containedIp));
+  }
+
+  @Test 
+  public void testRegexNotContained() throws UnknownHostException {
+	  assertFalse(whitelist(containedRegex).contains(notContainedIp));
+  }
+
 }