From 1b7990897e57df2d095e6eb06aa6c27bc81195ba Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 19 Aug 2012 00:14:16 +1200 Subject: [PATCH] Command-line options for header setting. --- libmproxy/cmdline.py | 105 +++++++++++++++++++++++++--------- libmproxy/console/__init__.py | 4 ++ libmproxy/dump.py | 5 ++ libmproxy/proxy.py | 5 ++ test/test_cmdline.py | 12 +++- 5 files changed, 102 insertions(+), 29 deletions(-) diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py index 57fce2c53..279c3cb4f 100644 --- a/libmproxy/cmdline.py +++ b/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) diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py index e835340e1..ee54afc2c 100644 --- a/libmproxy/console/__init__.py +++ b/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) diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 95e5ac0f7..98596b538 100644 --- a/libmproxy/dump.py +++ b/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), diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 09c56569a..a258b9418 100644 --- a/libmproxy/proxy.py +++ b/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] diff --git a/test/test_cmdline.py b/test/test_cmdline.py index 06b3abc9c..0a1d6e8fd 100644 --- a/test/test_cmdline.py +++ b/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 )