package us.kbase.narrativejobservice.db;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

public class SanitizeMongoObject {

	public static void sanitize(final Map<String, Object> data) {
		alterKeys(data, new KeyModifier() {
			
			@Override
			public String modify(String key) {
				//slightly inefficient, but meh
				return key.replace("%", "%25").replace("$", "%24")
						.replace(".", "%2e");
			}
		});
	}
	
	public static void befoul(final Map<String, Object> data) {
		alterKeys(data, new KeyModifier() {
			
			@Override
			public String modify(String key) {
				//slightly inefficient, but meh
				return key.replace("%2e", ".").replace("%24", "$")
						.replace("%25", "%");
			}
		});
	}
	
	private static interface KeyModifier {
		public String modify(String key);
	}

	//rewrite w/o recursion?
	private static Object alterKeys(final Object o, final KeyModifier mod) {
		if (o instanceof String || o instanceof Number ||
				o instanceof Boolean || o == null) {
			return o;
		} else if (o instanceof List) {
			@SuppressWarnings("unchecked")
			final List<Object> l = (List<Object>)o;
			for (Object lo: l) {
				alterKeys(lo, mod);
			}
			return o;
		} else if (o instanceof Map) {
			@SuppressWarnings("unchecked")
			final Map<String, Object> m = (Map<String, Object>)o;
			//save updated keys in separate map so we don't overwrite
			//keys before they're escaped
			final Map<String, Object> newm = new HashMap<String, Object>();
			final Iterator<Entry<String, Object>> iter = m.entrySet().iterator();
			while (iter.hasNext()) {
				final Entry<String, Object> e = iter.next();
				final Object value = alterKeys(e.getValue(), mod);
				final String newkey = mod.modify(e.getKey());
				iter.remove();
				newm.put(newkey, value);
			}
			m.putAll(newm);
			return o;
		} else {
			throw new IllegalStateException(
					"Unsupported class: " + o.getClass());
		}
	}
}

@@ -65,35 +65,8 @@
 ap;%0A
-import java.util.Iterator;%0A
 impo
@@ -1507,22 +1507,18 @@
 %09%09%09f
+or (f
 inal 
-Iterator%3C
 Entr
@@ -1538,16 +1538,11 @@
 ect%3E
-%3E iter =
+ e:
  m.e
@@ -1554,255 +1554,98 @@
 et()
-.iterator();%0A%09%09%09while (iter.hasNext()) %7B%0A%09%09%09%09final Entry%3CString, Object%3E e = iter.next();%0A%09%09%09%09final Object value = alterKeys(e.getValue(), mod);%0A%09%09%09%09final String newkey = mod.modify(e.getKey());%0A%09%09%09%09iter.remove();%0A%09%09%09%09newm.put(newkey, value);%0A%09%09%09%7D
+) %7B%0A%09%09%09%09newm.put(mod.modify(e.getKey()), alterKeys(e.getValue(), mod));%0A%09%09%09%7D%0A%09%09%09m.clear();
 %0A%09%09%09

Simplify SanitizeMongoObject method