summaryrefslogtreecommitdiff
path: root/crocoite/controller.py
diff options
context:
space:
mode:
Diffstat (limited to 'crocoite/controller.py')
-rw-r--r--crocoite/controller.py590
1 files changed, 453 insertions, 137 deletions
diff --git a/crocoite/controller.py b/crocoite/controller.py
index bc6f948..8374b4e 100644
--- a/crocoite/controller.py
+++ b/crocoite/controller.py
@@ -22,115 +22,297 @@
Controller classes, handling actions required for archival
"""
+import time, tempfile, asyncio, json, os, shutil, signal
+from itertools import islice
+from datetime import datetime
+from operator import attrgetter
+from abc import ABC, abstractmethod
+from yarl import URL
+
+from . import behavior as cbehavior
+from .browser import SiteLoader, RequestResponsePair, PageIdle, FrameNavigated
+from .util import getFormattedViewportMetrics, getSoftwareInfo
+from .behavior import ExtractLinksEvent
+from .devtools import toCookieParam
+
class ControllerSettings:
- def __init__ (self, logBuffer=1000, maxBodySize=50*1024*1024, idleTimeout=2, timeout=10):
- self.logBuffer = logBuffer
- self.maxBodySize = maxBodySize
+ __slots__ = ('idleTimeout', 'timeout', 'insecure', 'cookies')
+
+ def __init__ (self, idleTimeout=2, timeout=10, insecure=False, cookies=None):
self.idleTimeout = idleTimeout
self.timeout = timeout
+ self.insecure = insecure
+ self.cookies = cookies or []
- def toDict (self):
- return dict (logBuffer=self.logBuffer, maxBodySize=self.maxBodySize,
- idleTimeout=self.idleTimeout, timeout=self.timeout)
+ def __repr__ (self):
+ return f'<ControllerSetting idleTimeout={self.idleTimeout!r}, timeout={self.timeout!r}, insecure={self.insecure!r}, cookies={self.cookies!r}>'
defaultSettings = ControllerSettings ()
-import logging
-from urllib.parse import urlsplit, urlunsplit
+class EventHandler (ABC):
+ """ Abstract base class for event handler """
-import pychrome
+ __slots__ = ()
-from . import behavior as cbehavior
-from .browser import ChromeService
-from .warc import WarcLoader, SerializingWARCWriter
-from .util import getFormattedViewportMetrics
+ @abstractmethod
+ async def push (self, item):
+ raise NotImplementedError ()
+
+class StatsHandler (EventHandler):
+ __slots__ = ('stats', )
+
+ def __init__ (self):
+ self.stats = {'requests': 0, 'finished': 0, 'failed': 0, 'bytesRcv': 0}
+
+ async def push (self, item):
+ if isinstance (item, RequestResponsePair):
+ self.stats['requests'] += 1
+ if not item.response:
+ self.stats['failed'] += 1
+ else:
+ self.stats['finished'] += 1
+ self.stats['bytesRcv'] += item.response.bytesReceived
+
+class LogHandler (EventHandler):
+ """ Handle items by logging information about them """
+
+ __slots__ = ('logger', )
+
+ def __init__ (self, logger):
+ self.logger = logger.bind (context=type (self).__name__)
+
+ async def push (self, item):
+ if isinstance (item, ExtractLinksEvent):
+ # limit number of links per message, so json blob won’t get too big
+ it = iter (item.links)
+ limit = 100
+ while True:
+ limitlinks = list (islice (it, 0, limit))
+ if not limitlinks:
+ break
+ self.logger.info ('extracted links', context=type (item).__name__,
+ uuid='8ee5e9c9-1130-4c5c-88ff-718508546e0c', links=limitlinks)
+
+
+class ControllerStart:
+ __slots__ = ('payload', )
+
+ def __init__ (self, payload):
+ self.payload = payload
+
+class IdleStateTracker (EventHandler):
+ """ Track SiteLoader’s idle state by listening to PageIdle events """
+
+ __slots__ = ('_idle', '_loop', '_idleSince')
+
+ def __init__ (self, loop):
+ self._idle = True
+ self._loop = loop
+
+ self._idleSince = self._loop.time ()
+
+ async def push (self, item):
+ if isinstance (item, PageIdle):
+ self._idle = bool (item)
+ if self._idle:
+ self._idleSince = self._loop.time ()
+
+ async def wait (self, timeout):
+ """ Wait until page has been idle for at least timeout seconds. If the
+ page has been idle before calling this function it may return
+ immediately. """
+
+ assert timeout > 0
+ while True:
+ if self._idle:
+ now = self._loop.time ()
+ sleep = timeout-(now-self._idleSince)
+ if sleep <= 0:
+ break
+ else:
+ # not idle, check again after timeout expires
+ sleep = timeout
+ await asyncio.sleep (sleep)
+
+class InjectBehaviorOnload (EventHandler):
+ """ Control behavior script injection based on frame navigation messages.
+ When a page is reloaded (for whatever reason), the scripts need to be
+ reinjected. """
+
+ __slots__ = ('controller', '_loaded')
+
+ def __init__ (self, controller):
+ self.controller = controller
+ self._loaded = False
+
+ async def push (self, item):
+ if isinstance (item, FrameNavigated):
+ await self._runon ('load')
+ self._loaded = True
+
+ async def stop (self):
+ if self._loaded:
+ await self._runon ('stop')
-def firstOrNone (it):
- """ Return first item of iterator it or None if empty """
- try:
- return next (it)
- except StopIteration:
- return None
+ async def finish (self):
+ if self._loaded:
+ await self._runon ('finish')
+
+ async def _runon (self, method):
+ controller = self.controller
+ for b in controller._enabledBehavior:
+ f = getattr (b, 'on' + method)
+ async for item in f ():
+ await controller.processItem (item)
class SinglePageController:
"""
- Archive a single page url to file output.
+ Archive a single page url.
+
+ Dispatches between producer (site loader and behavior scripts) and consumer
+ (stats, warc writer).
"""
- def __init__ (self, url, output, service=ChromeService (), behavior=cbehavior.available, \
- logger=logging.getLogger(__name__), settings=defaultSettings):
+ __slots__ = ('url', 'service', 'behavior', 'settings', 'logger', 'handler',
+ 'warcinfo', '_enabledBehavior')
+
+ def __init__ (self, url, logger, \
+ service, behavior=cbehavior.available, \
+ settings=defaultSettings, handler=None, \
+ warcinfo=None):
self.url = url
- self.output = output
self.service = service
self.behavior = behavior
self.settings = settings
- self.logger = logger
-
- def run (self):
- ret = {'stats': None, 'links': []}
-
- with self.service as browser:
- browser = pychrome.Browser (url=browser)
- writer = SerializingWARCWriter (self.output, gzip=True)
-
- with WarcLoader (browser, self.url, writer,
- logBuffer=self.settings.logBuffer,
- maxBodySize=self.settings.maxBodySize) as l:
- version = l.tab.Browser.getVersion ()
- payload = {
- 'software': __package__,
- 'browser': version['product'],
+ self.logger = logger.bind (context=type (self).__name__, url=url)
+ self.handler = handler or []
+ self.warcinfo = warcinfo
+
+ async def processItem (self, item):
+ for h in self.handler:
+ await h.push (item)
+
+ async def run (self):
+ logger = self.logger
+ async def processQueue ():
+ async for item in l:
+ await self.processItem (item)
+
+ idle = IdleStateTracker (asyncio.get_event_loop ())
+ self.handler.append (idle)
+ behavior = InjectBehaviorOnload (self)
+ self.handler.append (behavior)
+
+ async with self.service as browser, SiteLoader (browser, logger=logger) as l:
+ handle = asyncio.ensure_future (processQueue ())
+ timeoutProc = asyncio.ensure_future (asyncio.sleep (self.settings.timeout))
+
+ # configure browser
+ tab = l.tab
+ await tab.Security.setIgnoreCertificateErrors (ignore=self.settings.insecure)
+ await tab.Network.setCookies (cookies=list (map (toCookieParam, self.settings.cookies)))
+
+ # not all behavior scripts are allowed for every URL, filter them
+ self._enabledBehavior = list (filter (lambda x: self.url in x,
+ map (lambda x: x (l, logger), self.behavior)))
+
+ version = await tab.Browser.getVersion ()
+ payload = {
+ 'software': getSoftwareInfo (),
+ 'browser': {
+ 'product': version['product'],
'useragent': version['userAgent'],
- 'viewport': getFormattedViewportMetrics (l.tab),
- }
- warcinfo = writer.create_warcinfo_record (filename=None, info=payload)
- writer.write_record (warcinfo)
-
- # not all behavior scripts are allowed for every URL, filter them
- enabledBehavior = list (filter (lambda x: self.url in x,
- map (lambda x: x (l), self.behavior)))
- linksBehavior = firstOrNone (filter (lambda x: isinstance (x, cbehavior.ExtractLinks),
- enabledBehavior))
-
- for b in enabledBehavior:
- self.logger.debug ('starting onload behavior {}'.format (b.name))
- b.onload ()
- l.start ()
-
- l.waitIdle (self.settings.idleTimeout, self.settings.timeout)
-
- for b in enabledBehavior:
- self.logger.debug ('starting onstop behavior {}'.format (b.name))
- b.onstop ()
-
- # if we stopped due to timeout, wait for remaining assets
- l.waitIdle (2, 60)
- l.stop ()
-
- for b in enabledBehavior:
- self.logger.debug ('starting onfinish behavior {}'.format (b.name))
- b.onfinish ()
-
- ret['stats'] = l.stats
- ret['links'] = linksBehavior.links if linksBehavior else None
- writer.flush ()
- return ret
-
-from collections import UserDict
-
-class IntegerDict (UserDict):
- """ Dict with dict/dict per-item arithmetic propagation, i.e. {1: 2}+{1: 1}={1: 3} """
- def __add__ (self, b):
- newdict = self.__class__ (self)
- for k, v in b.items ():
- if k in self:
- newdict[k] += v
+ 'viewport': await getFormattedViewportMetrics (tab),
+ },
+ 'tool': 'crocoite-single', # not the name of the cli utility
+ 'parameters': {
+ 'url': self.url,
+ 'idleTimeout': self.settings.idleTimeout,
+ 'timeout': self.settings.timeout,
+ 'behavior': list (map (attrgetter('name'), self._enabledBehavior)),
+ 'insecure': self.settings.insecure,
+ 'cookies': list (map (lambda x: x.OutputString(), self.settings.cookies)),
+ },
+ }
+ if self.warcinfo:
+ payload['extra'] = self.warcinfo
+ await self.processItem (ControllerStart (payload))
+
+ await l.navigate (self.url)
+
+ idleProc = asyncio.ensure_future (idle.wait (self.settings.idleTimeout))
+ while True:
+ try:
+ finished, pending = await asyncio.wait([idleProc, timeoutProc, handle],
+ return_when=asyncio.FIRST_COMPLETED)
+ except asyncio.CancelledError:
+ idleProc.cancel ()
+ timeoutProc.cancel ()
+ break
+
+ if handle in finished:
+ # something went wrong while processing the data
+ logger.error ('fetch failed',
+ uuid='43a0686a-a3a9-4214-9acd-43f6976f8ff3')
+ idleProc.cancel ()
+ timeoutProc.cancel ()
+ handle.result ()
+ assert False # previous line should always raise Exception
+ elif timeoutProc in finished:
+ # global timeout
+ logger.debug ('global timeout',
+ uuid='2f858adc-9448-4ace-94b4-7cd1484c0728')
+ idleProc.cancel ()
+ timeoutProc.result ()
+ break
+ elif idleProc in finished:
+ # idle timeout
+ logger.debug ('idle timeout',
+ uuid='90702590-94c4-44ef-9b37-02a16de444c3')
+ idleProc.result ()
+ timeoutProc.cancel ()
+ break
+
+ await behavior.stop ()
+ await tab.Page.stopLoading ()
+ await asyncio.sleep (1)
+ await behavior.finish ()
+
+ # wait until loads from behavior scripts are done and browser is
+ # idle for at least 1 second
+ try:
+ await asyncio.wait_for (idle.wait (1), timeout=1)
+ except (asyncio.TimeoutError, asyncio.CancelledError):
+ pass
+
+ if handle.done ():
+ handle.result ()
else:
- newdict[k] = v
- return newdict
+ handle.cancel ()
+
+class SetEntry:
+ """ A object, to be used with sets, that compares equality only on its
+ primary property. """
+ def __init__ (self, value, **props):
+ self.value = value
+ for k, v in props.items ():
+ setattr (self, k, v)
+
+ def __eq__ (self, b):
+ assert isinstance (b, SetEntry)
+ return self.value == b.value
+
+ def __hash__ (self):
+ return hash (self.value)
+
+ def __repr__ (self):
+ return f'<SetEntry {self.value!r}>'
class RecursionPolicy:
""" Abstract recursion policy """
+
+ __slots__ = ()
+
def __call__ (self, urls):
raise NotImplementedError
@@ -138,20 +320,23 @@ class DepthLimit (RecursionPolicy):
"""
Limit recursion by depth.
- depth==0 means no recursion, depth==1 is the page and outgoing links, …
+ depth==0 means no recursion, depth==1 is the page and outgoing links
"""
+
+ __slots__ = ('maxdepth', )
+
def __init__ (self, maxdepth=0):
self.maxdepth = maxdepth
def __call__ (self, urls):
- if self.maxdepth <= 0:
- return {}
- else:
- self.maxdepth -= 1
- return urls
+ newurls = set ()
+ for u in urls:
+ if u.depth <= self.maxdepth:
+ newurls.add (u)
+ return newurls
def __repr__ (self):
- return '<DepthLimit {}>'.format (self.maxdepth)
+ return f'<DepthLimit {self.maxdepth}>'
class PrefixLimit (RecursionPolicy):
"""
@@ -161,65 +346,196 @@ class PrefixLimit (RecursionPolicy):
ignored: http://example.com/bar http://offsite.example/foo
accepted: http://example.com/foobar http://example.com/foo/bar
"""
+
+ __slots__ = ('prefix', )
+
def __init__ (self, prefix):
self.prefix = prefix
def __call__ (self, urls):
- return set (filter (lambda u: u.startswith (self.prefix), urls))
+ return set (filter (lambda u: str(u.value).startswith (str (self.prefix)), urls))
-def removeFragment (u):
- """ Remove fragment from url (i.e. #hashvalue) """
- s = urlsplit (u)
- return urlunsplit ((s.scheme, s.netloc, s.path, s.query, ''))
+def hasTemplate (s):
+ """ Return True if string s has string templates """
+ return '{' in s and '}' in s
class RecursiveController:
"""
Simple recursive controller
- Visits links acording to recursionPolicy
+ Visits links acording to policy
"""
- def __init__ (self, url, output, service=ChromeService (), behavior=cbehavior.available, \
- logger=logging.getLogger(__name__), settings=defaultSettings,
- recursionPolicy=DepthLimit (0)):
+ __slots__ = ('url', 'output', 'command', 'logger', 'policy', 'have',
+ 'pending', 'stats', 'tempdir', 'running', 'concurrency',
+ 'copyLock')
+
+ SCHEME_WHITELIST = {'http', 'https'}
+
+ def __init__ (self, url, output, command, logger,
+ tempdir=None, policy=DepthLimit (0), concurrency=1):
self.url = url
self.output = output
- self.service = service
- self.behavior = behavior
- self.settings = settings
- self.logger = logger
- self.recursionPolicy = recursionPolicy
-
- def fetch (self, urls):
+ self.command = command
+ self.logger = logger.bind (context=type(self).__name__, seedurl=url)
+ self.policy = policy
+ self.tempdir = tempdir
+ # A lock if only a single output file (no template) is requested
+ self.copyLock = None if hasTemplate (output) else asyncio.Lock ()
+ # some sanity checks. XXX move to argparse?
+ if self.copyLock and os.path.exists (self.output):
+ raise ValueError ('Output file exists')
+ # tasks currently running
+ self.running = set ()
+ # max number of tasks running
+ self.concurrency = concurrency
+ # keep in sync with StatsHandler
+ self.stats = {'requests': 0, 'finished': 0, 'failed': 0, 'bytesRcv': 0, 'crashed': 0, 'ignored': 0}
+
+ async def fetch (self, entry, seqnum):
"""
- Overrideable fetch action for URLs. Defaults to sequential
- SinglePageController.
+ Fetch a single URL using an external command
+
+ command is usually crocoite-single
"""
- result = []
- for u in urls:
- c = SinglePageController (u, self.output, self.service,
- self.behavior, self.logger, self.settings)
- result.append (c.run ())
- return result
-
- def run (self):
- have = set ()
- urls = set ([self.url])
- ret = {'stats': IntegerDict ()}
-
- while urls:
- self.logger.info ('retrieving {} urls'.format (len (urls)))
- result = self.fetch (urls)
-
- have.update (urls)
- urls = set ()
- for r in result:
- ret['stats'] += r['stats']
- urls.update (map (removeFragment, r['links']))
- urls.difference_update (have)
-
- urls = self.recursionPolicy (urls)
- # everything in ret must be serializeable
- ret['stats'] = dict (ret['stats'])
- return ret
+
+ assert isinstance (entry, SetEntry)
+
+ url = entry.value
+ depth = entry.depth
+ logger = self.logger.bind (url=url)
+
+ def formatCommand (e):
+ # provide means to disable variable expansion
+ if e.startswith ('!'):
+ return e[1:]
+ else:
+ return e.format (url=url, dest=dest.name)
+
+ def formatOutput (p):
+ return p.format (host=url.host,
+ date=datetime.utcnow ().isoformat (), seqnum=seqnum)
+
+ def logStats ():
+ logger.info ('stats', uuid='24d92d16-770e-4088-b769-4020e127a7ff', **self.stats)
+
+ if url.scheme not in self.SCHEME_WHITELIST:
+ self.stats['ignored'] += 1
+ logStats ()
+ self.logger.warning ('scheme not whitelisted', url=url,
+ uuid='57e838de-4494-4316-ae98-cd3a2ebf541b')
+ return
+
+ dest = tempfile.NamedTemporaryFile (dir=self.tempdir,
+ prefix=os.path.basename (self.output) + '-', suffix='.warc.gz',
+ delete=False)
+ command = list (map (formatCommand, self.command))
+ logger.info ('fetch', uuid='d1288fbe-8bae-42c8-af8c-f2fa8b41794f',
+ command=command)
+ try:
+ process = await asyncio.create_subprocess_exec (*command,
+ stdout=asyncio.subprocess.PIPE,
+ stderr=asyncio.subprocess.DEVNULL,
+ stdin=asyncio.subprocess.DEVNULL,
+ start_new_session=True, limit=100*1024*1024)
+ while True:
+ data = await process.stdout.readline ()
+ if not data:
+ break
+ data = json.loads (data)
+ uuid = data.get ('uuid')
+ if uuid == '8ee5e9c9-1130-4c5c-88ff-718508546e0c':
+ links = set (self.policy (map (lambda x: SetEntry (URL(x).with_fragment(None), depth=depth+1), data.get ('links', []))))
+ links.difference_update (self.have)
+ self.pending.update (links)
+ elif uuid == '24d92d16-770e-4088-b769-4020e127a7ff':
+ for k in self.stats.keys ():
+ self.stats[k] += data.get (k, 0)
+ logStats ()
+ except asyncio.CancelledError:
+ # graceful cancellation
+ process.send_signal (signal.SIGINT)
+ except Exception as e:
+ process.kill ()
+ raise e
+ finally:
+ code = await process.wait()
+ if code == 0:
+ if self.copyLock is None:
+ # atomically move once finished
+ lastDestpath = None
+ while True:
+ # XXX: must generate a new name every time, otherwise
+ # this loop never terminates
+ destpath = formatOutput (self.output)
+ assert destpath != lastDestpath
+ lastDestpath = destpath
+
+ # python does not have rename(…, …, RENAME_NOREPLACE),
+ # but this is safe nontheless, since we’re
+ # single-threaded
+ if not os.path.exists (destpath):
+ # create the directory, so templates like
+ # /{host}/{date}/… are possible
+ os.makedirs (os.path.dirname (destpath), exist_ok=True)
+ os.rename (dest.name, destpath)
+ break
+ else:
+ # atomically (in the context of this process) append to
+ # existing file
+ async with self.copyLock:
+ with open (dest.name, 'rb') as infd, \
+ open (self.output, 'ab') as outfd:
+ shutil.copyfileobj (infd, outfd)
+ os.unlink (dest.name)
+ else:
+ self.stats['crashed'] += 1
+ logStats ()
+
+ async def run (self):
+ def log ():
+ # self.have includes running jobs
+ self.logger.info ('recursing',
+ uuid='5b8498e4-868d-413c-a67e-004516b8452c',
+ pending=len (self.pending),
+ have=len (self.have)-len(self.running),
+ running=len (self.running))
+
+ seqnum = 1
+ try:
+ self.have = set ()
+ self.pending = set ([SetEntry (self.url, depth=0)])
+
+ while self.pending:
+ # since pending is a set this picks a random item, which is fine
+ u = self.pending.pop ()
+ self.have.add (u)
+ t = asyncio.ensure_future (self.fetch (u, seqnum))
+ self.running.add (t)
+ seqnum += 1
+
+ log ()
+
+ if len (self.running) >= self.concurrency or not self.pending:
+ done, pending = await asyncio.wait (self.running,
+ return_when=asyncio.FIRST_COMPLETED)
+ self.running.difference_update (done)
+ # propagate exceptions
+ for r in done:
+ r.result ()
+ except asyncio.CancelledError:
+ self.logger.info ('cancel',
+ uuid='d58154c8-ec27-40f2-ab9e-e25c1b21cd88',
+ pending=len (self.pending),
+ have=len (self.have)-len (self.running),
+ running=len (self.running))
+ finally:
+ done = await asyncio.gather (*self.running,
+ return_exceptions=True)
+ # propagate exceptions
+ for r in done:
+ if isinstance (r, Exception):
+ raise r
+ self.running = set ()
+ log ()
generated by cgit v1.2.3 (git 2.39.1) at 2024-05-01 23:03:15 +0000