summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--crocoite/test_tools.py188
-rw-r--r--crocoite/tools.py35
2 files changed, 205 insertions, 18 deletions
diff --git a/crocoite/test_tools.py b/crocoite/test_tools.py
new file mode 100644
index 0000000..09886b2
--- /dev/null
+++ b/crocoite/test_tools.py
@@ -0,0 +1,188 @@
+# Copyright (c) 2018 crocoite contributors
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from tempfile import NamedTemporaryFile
+from operator import itemgetter
+from io import BytesIO
+import pytest
+from warcio.archiveiterator import ArchiveIterator
+from warcio.warcwriter import WARCWriter
+from warcio.statusandheaders import StatusAndHeaders
+
+from .tools import mergeWarc
+
+@pytest.fixture
+def writer():
+ return WARCWriter (NamedTemporaryFile(), gzip=True)
+
+def recordsEqual(golden, underTest):
+ for a, b in zip (golden, underTest):
+ # record ids are not predictable, so we cannot compare them
+ a.rec_headers.remove_header('WARC-Record-Id')
+ a.rec_headers.remove_header('WARC-Block-Digest')
+ b.rec_headers.remove_header('WARC-Record-Id')
+ b.rec_headers.remove_header('WARC-Block-Digest')
+ aheader = sorted(a.rec_headers.headers, key=itemgetter(0))
+ bheader = sorted(b.rec_headers.headers, key=itemgetter(0))
+ assert aheader == bheader
+
+def test_unmodified(writer):
+ """
+ Single request/response pair, no revisits
+ """
+
+ records = []
+ httpHeaders = StatusAndHeaders('GET / HTTP/1.1', {}, is_http_request=True)
+ warcHeaders = {}
+ record = writer.create_warc_record ('http://example.com/', 'request', payload=BytesIO(b'foobar'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ httpHeaders = StatusAndHeaders('200 OK', {}, protocol='HTTP/1.1')
+ record = writer.create_warc_record ('http://example.com/', 'response', payload=BytesIO(b'data'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ for r in records:
+ writer.write_record (r)
+
+ output = NamedTemporaryFile()
+ mergeWarc ([writer.out.name], output)
+
+ output.seek(0)
+ recordsEqual (records, ArchiveIterator (output))
+
+def test_different_payload(writer):
+ """
+ Duplicate URL, but different payload
+ """
+
+ records = []
+ for i in range (2):
+ httpHeaders = StatusAndHeaders('GET / HTTP/1.1', {}, is_http_request=True)
+ warcHeaders = {}
+ record = writer.create_warc_record ('http://example.com/', 'request', payload=BytesIO(b'foobar'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ httpHeaders = StatusAndHeaders('200 OK', {}, protocol='HTTP/1.1')
+ record = writer.create_warc_record ('http://example.com/', 'response',
+ payload=BytesIO('data{}'.format(i).encode ('utf8')),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ for r in records:
+ writer.write_record (r)
+
+ output = NamedTemporaryFile()
+ mergeWarc ([writer.out.name], output)
+
+ output.seek(0)
+ recordsEqual (records, ArchiveIterator (output))
+
+def makeRevisit(writer, ref, dup):
+ """ Make revisit record for reference """
+ dupHeaders = dup.rec_headers
+ refHeaders = ref.rec_headers
+ httpHeaders = StatusAndHeaders('200 OK', {}, protocol='HTTP/1.1')
+ record = writer.create_revisit_record (dupHeaders.get_header('WARC-Target-URI'),
+ digest=refHeaders.get_header('WARC-Payload-Digest'),
+ refers_to_uri=refHeaders.get_header('WARC-Target-URI'),
+ refers_to_date=refHeaders.get_header('WARC-Date'))
+ record.rec_headers.add_header ('WARC-Refers-To', refHeaders.get_header('WARC-Record-ID'))
+ record.rec_headers.add_header ('WARC-Truncated', 'length')
+ record.rec_headers.add_header ('Content-Length', '0')
+ # XXX: added by warcio, but this seems wrong
+ record.rec_headers.add_header ('Content-Type', 'application/http; msgtype=response')
+ return record
+
+def test_resp_revisit_same_url(writer):
+ """
+ Duplicate record for the same URL, creates a revisit
+ """
+
+ records = []
+ for i in range (2):
+ httpHeaders = StatusAndHeaders('GET / HTTP/1.1', {}, is_http_request=True)
+ warcHeaders = {}
+ record = writer.create_warc_record ('http://example.com/', 'request', payload=BytesIO(b'foobar'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ httpHeaders = StatusAndHeaders('200 OK', {}, protocol='HTTP/1.1')
+ record = writer.create_warc_record ('http://example.com/', 'response', payload=BytesIO(b'data'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ for r in records:
+ writer.write_record (r)
+
+ dup = records.pop ()
+ ref = records[1]
+ records.append (makeRevisit (writer, ref, dup))
+
+ output = NamedTemporaryFile()
+ mergeWarc ([writer.out.name], output)
+
+ output.seek(0)
+ recordsEqual (records, ArchiveIterator (output))
+
+def test_resp_revisit_other_url(writer):
+ """
+ Duplicate record for different URL, creates a revisit
+ """
+
+ records = []
+
+ httpHeaders = StatusAndHeaders('GET / HTTP/1.1', {}, is_http_request=True)
+ warcHeaders = {}
+ record = writer.create_warc_record ('http://example.com/', 'request', payload=BytesIO(b'foobar'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ httpHeaders = StatusAndHeaders('200 OK', {}, protocol='HTTP/1.1')
+ record = writer.create_warc_record ('http://example.com/', 'response', payload=BytesIO(b'data'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ httpHeaders = StatusAndHeaders('GET / HTTP/1.1', {}, is_http_request=True)
+ warcHeaders = {}
+ record = writer.create_warc_record ('http://example.com/one', 'request', payload=BytesIO(b'foobar'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ httpHeaders = StatusAndHeaders('200 OK', {}, protocol='HTTP/1.1')
+ record = writer.create_warc_record ('http://example.com/one', 'response', payload=BytesIO(b'data'),
+ warc_headers_dict=warcHeaders, http_headers=httpHeaders)
+ records.append (record)
+
+ for r in records:
+ writer.write_record (r)
+
+ dup = records.pop ()
+ ref = records[1]
+ records.append (makeRevisit (writer, ref, dup))
+
+ output = NamedTemporaryFile()
+ mergeWarc ([writer.out.name], output)
+
+ output.seek(0)
+ recordsEqual (records, ArchiveIterator (output))
+
diff --git a/crocoite/tools.py b/crocoite/tools.py
index 3aeaaad..8541ca2 100644
--- a/crocoite/tools.py
+++ b/crocoite/tools.py
@@ -26,26 +26,12 @@ import shutil, sys, re, os, logging, argparse
from warcio.archiveiterator import ArchiveIterator
from warcio.warcwriter import WARCWriter
-def mergeWarc ():
- """
- Merge multiple WARC files into a single file, writing revisit records for
- items which occur multiple times
- """
-
- parser = argparse.ArgumentParser(description='Merge WARCs, reads filenames from stdin.')
- parser.add_argument('--verbose', '-v', action='store_true')
- parser.add_argument('output', type=argparse.FileType ('wb'), help='Output WARC')
-
- args = parser.parse_args()
- loglevel = logging.DEBUG if args.verbose else logging.INFO
- logging.basicConfig (level=loglevel)
-
+def mergeWarc (files, output):
unique = 0
revisit = 0
payloadMap = {}
- writer = WARCWriter (args.output, gzip=True)
- for l in sys.stdin:
- l = l.strip ()
+ writer = WARCWriter (output, gzip=True)
+ for l in files:
with open (l, 'rb') as fd:
for record in ArchiveIterator (fd):
if record.rec_type in {'resource', 'response'}:
@@ -59,7 +45,9 @@ def mergeWarc ():
unique += 1
else:
logging.debug ('Record {} is duplicate of {}'.format (rid, dup['id']))
- record = writer.create_revisit_record (dup['uri'], csum, dup['uri'], dup['date'])
+ record = writer.create_revisit_record (
+ headers.get_header('WARC-Target-URI'), digest=csum,
+ refers_to_uri=dup['uri'], refers_to_date=dup['date'])
record.rec_headers.add_header ('WARC-Truncated', 'length')
record.rec_headers.add_header ('WARC-Refers-To', dup['id'])
revisit += 1
@@ -68,6 +56,17 @@ def mergeWarc ():
writer.write_record (record)
logging.info ('Wrote {} unique records, {} revisits'.format (unique, revisit))
+def mergeWarcCli():
+ parser = argparse.ArgumentParser(description='Merge WARCs, reads filenames from stdin.')
+ parser.add_argument('--verbose', '-v', action='store_true')
+ parser.add_argument('output', type=argparse.FileType ('wb'), help='Output WARC')
+
+ args = parser.parse_args()
+ loglevel = logging.DEBUG if args.verbose else logging.INFO
+ logging.basicConfig (level=loglevel)
+
+ mergeWarc([l.strip() for l in sys.stdin], args.output)
+
def extractScreenshot ():
"""
Extract page screenshots from a WARC generated by crocoite into files