Add -c and -C options to mitmdump to control sticky cookies.

It's dumb that this needs to be two different options, but optparse doesn't
support optional arguments. It would be much nicer to just have "-c" for "all",
and "-c filter" for a specified filter.

3 files changed, 24 insertions, 4 deletions

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 396dc31d..c1f89d88 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -13,6 +13,7 @@ class Options(object):
         "verbosity",
         "wfile",
         "rheaders",
+        "stickycookie",
     ]
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -38,6 +39,9 @@ class DumpMaster(flow.FlowMaster):
         if self.o.request_script:
             self.set_request_script(self.o.request_script)  
 
+        if options.stickycookie:
+            self.set_stickycookie(options.stickycookie)
+
         if options.wfile:
             path = os.path.expanduser(options.wfile)
             try:
@@ -55,6 +59,7 @@ class DumpMaster(flow.FlowMaster):
                 raise DumpError(v.strerror)
             self.start_playback(flows, options.kill, options.rheaders)
 
+
     def _runscript(self, f, script):
         try:
             ret = f.run_script(script)
diff --git a/mitmdump b/mitmdump
index 66cb5ade..3560b191 100755
--- a/mitmdump
+++ b/mitmdump
@@ -18,6 +18,8 @@
 import sys, os.path
 from libmproxy import proxy, dump, utils
 from libmproxy import VERSION
+
+
 from optparse import OptionParser, OptionGroup
 
 
@@ -32,13 +34,19 @@ if __name__ == '__main__':
         type = "int", dest="port", default=8080,
         help = "Port."
     )
+    parser.add_option("-c", 
+                      action="store_true", dest="stickycookie_all", default=None,
+                      help="Set sticky cookie for all requests.")
+    parser.add_option("-C",
+                      action="store", dest="stickycookie_filt", default=None, metavar="FILTER",
+                      help="Set sticky cookie filter. Matched against requests.")
     parser.add_option("-q", 
                       action="store_true", dest="quiet",
                       help="Quiet.")
-    parser.add_option("", "--reqscript",
+    parser.add_option("--reqscript",
                       action="store", dest="request_script", default=None,
                       help="Script to run when a request is recieved.")
-    parser.add_option("", "--respscript",
+    parser.add_option("--respscript",
                       action="store", dest="response_script", default=None,
                       help="Script to run when a response is recieved.")
     parser.add_option("-v", 
@@ -65,10 +73,14 @@ if __name__ == '__main__':
 
     options, args = parser.parse_args()
 
-
     if options.quiet:
         options.verbose = 0
 
+    stickycookie = None
+    if options.stickycookie_all:
+        stickycookie = ".*"
+    elif options.stickycookie_filt:
+        stickycookie = stickycookie_filt
 
     config = proxy.process_certificate_option_group(parser, options)
     server = proxy.ProxyServer(config, options.port)
@@ -79,7 +91,8 @@ if __name__ == '__main__':
         response_script = options.response_script,
         replay = options.replay,
         kill = options.kill,
-        rheaders = options.rheaders
+        rheaders = options.rheaders,
+        stickycookie = stickycookie
     )
     if args:
         filt = " ".join(args)
diff --git a/todo b/todo
index d9dca4fc..db47735c 100644
--- a/todo
+++ b/todo
@@ -4,6 +4,8 @@ Futures:
     - Post and URL field parsing and editing.
     - Upstream proxies.
     - Strings view for binary responses.
+    - Better manipulation of flows in mitmproxy: copying flows, adding
+    comments.
 
 Bugs: