Command-line options for header setting.

libmproxy/cmdline.py

@@ -17,9 +17,28 @@ import proxy
 import re, filt
 
 
-class ParseReplaceException(Exception): pass
+class ParseException(Exception): pass
 class OptionException(Exception): pass
 
+def _parse_hook(s):
+    sep, rem = s[0], s[1:]
+    parts = rem.split(sep, 2)
+    if len(parts) == 2:
+        patt = ".*"
+        a, b = parts
+    elif len(parts) == 3:
+        patt, a, b = parts
+    else:
+        raise ParseException("Malformed hook specifier - too few clauses: %s"%s)
+
+    if not a:
+        raise ParseException("Empty clause: %s"%str(patt))
+
+    if not filt.parse(patt):
+        raise ParseException("Malformed filter pattern: %s"%patt)
+
+    return patt, a, b
+
 
 def parse_replace_hook(s):
     """
@@ -45,32 +64,45 @@ def parse_replace_hook(s):
             /one/two/foo/bar/
 
         Checks that pattern and regex are both well-formed. Raises
-        ParseReplaceException on error.
+        ParseException on error.
     """
-    sep, rem = s[0], s[1:]
-    parts = rem.split(sep, 2)
-    if len(parts) == 2:
-        patt = ".*"
-        regex, replacement = parts
-    elif len(parts) == 3:
-        patt, regex, replacement = parts
-    else:
-        raise ParseReplaceException("Malformed replacement specifier - too few clauses: %s"%s)
-
-    if not regex:
-        raise ParseReplaceException("Empty replacement regex: %s"%str(patt))
-
+    patt, regex, replacement = _parse_hook(s)
     try:
         re.compile(regex)
     except re.error, e:
-        raise ParseReplaceException("Malformed replacement regex: %s"%str(e.message))
-
-    if not filt.parse(patt):
-        raise ParseReplaceException("Malformed replacement filter pattern: %s"%patt)
-
+        raise ParseException("Malformed replacement regex: %s"%str(e.message))
     return patt, regex, replacement
 
 
+def parse_setheader(s):
+    """
+        Returns a (pattern, header, value) tuple.
+
+        The general form for a replacement hook is as follows:
+
+            /patt/header/value
+
+        The first character specifies the separator. Example:
+
+            :~q:foo:bar
+
+        If only two clauses are specified, the pattern is set to match
+        universally (i.e. ".*"). Example:
+
+            /foo/bar/
+
+        Clauses are parsed from left to right. Extra separators are taken to be
+        part of the final clause. For instance, the value clause below is
+        "foo/bar/":
+
+            /one/two/foo/bar/
+
+        Checks that pattern and regex are both well-formed. Raises
+        ParseException on error.
+    """
+    return _parse_hook(s)
+
+
 def get_common_options(options):
     stickycookie, stickyauth = None, None
     if options.stickycookie_filt:
@@ -83,13 +115,13 @@ def get_common_options(options):
     for i in options.replace:
         try:
             p = parse_replace_hook(i)
-        except ParseReplaceException, e:
+        except ParseException, e:
             raise OptionException(e.message)
         reps.append(p)
     for i in options.replace_file:
         try:
             patt, rex, path = parse_replace_hook(i)
-        except ParseReplaceException, e:
+        except ParseException, e:
             raise OptionException(e.message)
         try:
             v = open(path, "r").read()
@@ -97,6 +129,15 @@ def get_common_options(options):
             raise OptionException("Could not read replace file: %s"%path)
         reps.append((patt, rex, v))
 
+
+    setheaders = []
+    for i in options.setheader:
+        try:
+            p = parse_setheader(i)
+        except ParseException, e:
+            raise OptionException(e.message)
+        setheaders.append(p)
+
     return dict(
         anticache = options.anticache,
         anticomp = options.anticomp,
@@ -108,6 +149,7 @@ def get_common_options(options):
         rheaders = options.rheaders,
         rfile = options.rfile,
         replacements = reps,
+        setheaders = setheaders,
         server_replay = options.server_replay,
         script = options.script,
         stickycookie = stickycookie,
@@ -211,6 +253,7 @@ def common_options(parser):
         action="store", dest="cert_wait_time", default=0,
         help="Wait for specified number of seconds after a new cert is generated. This can smooth over small discrepancies between the client and server times."
     )
+
     parser.add_argument(
         "--no-upstream-cert", default=False,
         action="store_true", dest="no_upstream_cert",
@@ -271,14 +314,24 @@ def common_options(parser):
     group.add_argument(
         "--replace-from-file",
         action="append", type=str, dest="replace_file", default=[],
-        metavar="PATTERN",
+        metavar="PATH",
         help="Replacement pattern, where the replacement clause is a path to a file."
     )
 
+
+    group = parser.add_argument_group(
+        "Set Headers",
+        """
+            Header specifications are of the form "/pattern/header/value",
+            where the separator can be any character. Please see the
+            documentation for more information.
+        """.strip()
+    )
     group.add_argument(
-        "--dummy-certs", action="store",
-        type = str, dest = "certdir", default=None,
-        help = "Generated dummy certs directory."
+        "--setheader",
+        action="append", type=str, dest="setheader", default=[],
+        metavar="PATTERN",
+        help="Header set pattern."
     )
 
     proxy.certificate_option_group(parser)

libmproxy/console/__init__.py

@@ -339,6 +339,7 @@ class Options(object):
         "script",
         "replacements",
         "rheaders",
+        "setheaders",
         "server_replay",
         "stickycookie",
         "stickyauth",
@@ -369,6 +370,9 @@ class ConsoleMaster(flow.FlowMaster):
         for i in options.replacements:
             self.replacehooks.add(*i)
 
+        for i in options.setheaders:
+            self.setheaders.add(*i)
+
         self.flow_list_walker = None
         self.set_palette(options.palette)
 

libmproxy/dump.py

@@ -33,6 +33,7 @@ class Options(object):
         "replacements",
         "rfile",
         "rheaders",
+        "setheaders",
         "server_replay",
         "script",
         "stickycookie",
@@ -99,6 +100,10 @@ class DumpMaster(flow.FlowMaster):
             for i in options.replacements:
                 self.replacehooks.add(*i)
 
+        if options.setheaders:
+            for i in options.setheaders:
+                self.setheaders.add(*i)
+
         if options.server_replay:
             self.start_server_playback(
                 self._readflow(options.server_replay),

libmproxy/proxy.py

@@ -475,6 +475,11 @@ def certificate_option_group(parser):
         type = str, dest = "clientcerts", default=None,
         help = "Client certificate directory."
     )
+    group.add_argument(
+        "--dummy-certs", action="store",
+        type = str, dest = "certdir", default=None,
+        help = "Generated dummy certs directory."
+    )
 
 
 TRANSPARENT_SSL_PORTS = [443, 8443]

test/test_cmdline.py

@@ -14,7 +14,7 @@ def test_parse_replace_hook():
     assert x == (".*", "bar", "voing")
 
     tutils.raises(
-        cmdline.ParseReplaceException,
+        cmdline.ParseException,
         cmdline.parse_replace_hook,
         "/foo"
     )
@@ -29,11 +29,17 @@ def test_parse_replace_hook():
         "/~/foo/rep"
     )
     tutils.raises(
-        "empty replacement regex",
+        "empty clause",
         cmdline.parse_replace_hook,
         "//"
     )
 
+
+def test_parse_setheaders():
+    x = cmdline.parse_replace_hook("/foo/bar/voing")
+    assert x == ("foo", "bar", "voing")
+
+
 def test_common():
     parser = argparse.ArgumentParser()
     cmdline.common_options(parser)
@@ -53,7 +59,7 @@ def test_common():
 
     opts.replace = ["//"]
     tutils.raises(
-        "empty replacement regex",
+        "empty clause",
         cmdline.get_common_options,
         opts
     )