From b54e4c1c06cabce17672b7ead53088d5eb007ee9 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Mon, 22 Dec 2014 20:34:18 -0800
Subject: [PATCH] tests: add more tests for cookie, html and rewrite_live crsf

---
 pywb/rewrite/test/test_cookie_rewriter.py | 22 +++++++++++++++++-----
 pywb/rewrite/test/test_html_rewriter.py   |  7 +++++++
 pywb/rewrite/test/test_rewrite_live.py    |  7 +++++++
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/pywb/rewrite/test/test_cookie_rewriter.py b/pywb/rewrite/test/test_cookie_rewriter.py
index 4f57464f..1984fb5a 100644
--- a/pywb/rewrite/test/test_cookie_rewriter.py
+++ b/pywb/rewrite/test/test_cookie_rewriter.py
@@ -24,17 +24,28 @@ r"""
 >>> rewrite_cookie('abc@def=123')
 []
 
-# ExactCookieRewriter
->>> rewrite_cookie('some=value; Path=/diff/path/;', urlrewriter, ExactPathCookieRewriter)
+
+# ExactCookieRewriter -- always removes Path and Domain
+>>> rewrite_cookie('some=value; Path=/diff/path/;', urlrewriter, 'exact')
 [('Set-Cookie', 'some=value')]
 
->>> rewrite_cookie('some=value; Domain=.example.com; Path=/diff/path/; Max-Age=1500', urlrewriter, ExactPathCookieRewriter)
+>>> rewrite_cookie('some=value; Domain=.example.com; Path=/diff/path/; Max-Age=1500', urlrewriter, 'exact')
 [('Set-Cookie', 'some=value')]
 
+
+# RootCookieRewriter -- always sets Path=/, removes Domain
+>>> rewrite_cookie('some=value; Path=/diff/path/;', urlrewriter, 'root')
+[('Set-Cookie', 'some=value; Path=/')]
+
+>>> rewrite_cookie('some=value; Domain=.example.com; Path=/diff/path/; Max-Age=1500', urlrewriter, 'root')
+[('Set-Cookie', 'some=value; Path=/')]
+
+
+
 """
 
 
-from pywb.rewrite.cookie_rewriter import MinimalScopeCookieRewriter, ExactPathCookieRewriter
+from pywb.rewrite.cookie_rewriter import MinimalScopeCookieRewriter, get_cookie_rewriter
 from pywb.rewrite.url_rewriter import UrlRewriter
 
 urlrewriter = UrlRewriter('20131226101010/http://example.com/some/path/index.html', '/pywb/')
@@ -42,6 +53,7 @@ urlrewriter = UrlRewriter('20131226101010/http://example.com/some/path/index.htm
 urlrewriter2 = UrlRewriter('em_/http://example.com/', '/preview/')
 
 
-def rewrite_cookie(cookie_str, rewriter=urlrewriter, cookie_rewriter=MinimalScopeCookieRewriter):
+def rewrite_cookie(cookie_str, rewriter=urlrewriter, scope='default'):
+    cookie_rewriter = get_cookie_rewriter(scope)
     return cookie_rewriter(rewriter).rewrite(cookie_str)
 
diff --git a/pywb/rewrite/test/test_html_rewriter.py b/pywb/rewrite/test/test_html_rewriter.py
index 31515250..7e17e28f 100644
--- a/pywb/rewrite/test/test_html_rewriter.py
+++ b/pywb/rewrite/test/test_html_rewriter.py
@@ -72,6 +72,13 @@ ur"""
 >>> parse('<div data-url="http://example.com/a/b/c.html" data-some-other-value="http://example.com/img.gif">')
 <div data-url="/web/20131226101010oe_/http://example.com/a/b/c.html" data-some-other-value="/web/20131226101010oe_/http://example.com/img.gif">
 
+# param tag -- rewrite conditionally if url
+>>> parse('<param value="http://example.com/"/>')
+<param value="/web/20131226101010oe_/http://example.com/"/>
+
+>>> parse('<param value="foo bar"/>')
+<param value="foo bar"/>
+
 # srcset attrib
 >>> parse('<img srcset="//example.com/1x 1x, //example.com/foo 2x, https://example.com/bar 4x">')
 <img srcset="/web/20131226101010/http://example.com/1x 1x, /web/20131226101010/http://example.com/foo 2x, /web/20131226101010/https://example.com/bar 4x">
diff --git a/pywb/rewrite/test/test_rewrite_live.py b/pywb/rewrite/test/test_rewrite_live.py
index 3ea189a4..fdea7557 100644
--- a/pywb/rewrite/test/test_rewrite_live.py
+++ b/pywb/rewrite/test/test_rewrite_live.py
@@ -18,6 +18,13 @@ def head_insert_func(rule, cdx):
     else:
         return ''
 
+def test_csrf_token_headers():
+    rewriter = LiveRewriter()
+    env = {'HTTP_X_CSRFTOKEN': 'wrong', 'HTTP_COOKIE': 'csrftoken=foobar'}
+
+    req_headers = rewriter.translate_headers('http://example.com/', env)
+
+    assert req_headers == {'X-CSRFToken': 'foobar', 'Cookie': 'csrftoken=foobar'}
 
 def test_local_1():
     status_headers, buff = get_rewritten(get_test_dir() + 'text_content/sample.html',