diff options
Diffstat (limited to 'crocoite/controller.py')
| -rw-r--r-- | crocoite/controller.py | 434 |
1 files changed, 289 insertions, 145 deletions
diff --git a/crocoite/controller.py b/crocoite/controller.py index f8b1420..8374b4e 100644 --- a/crocoite/controller.py +++ b/crocoite/controller.py @@ -22,58 +22,56 @@ Controller classes, handling actions required for archival """ -import time -import tempfile, asyncio, json, os +import time, tempfile, asyncio, json, os, shutil, signal from itertools import islice from datetime import datetime -from urllib.parse import urlparse from operator import attrgetter +from abc import ABC, abstractmethod +from yarl import URL from . import behavior as cbehavior -from .browser import SiteLoader, Item -from .util import getFormattedViewportMetrics, getSoftwareInfo, removeFragment +from .browser import SiteLoader, RequestResponsePair, PageIdle, FrameNavigated +from .util import getFormattedViewportMetrics, getSoftwareInfo from .behavior import ExtractLinksEvent +from .devtools import toCookieParam class ControllerSettings: - __slots__ = ('idleTimeout', 'timeout') + __slots__ = ('idleTimeout', 'timeout', 'insecure', 'cookies') - def __init__ (self, idleTimeout=2, timeout=10): + 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 (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 () -class EventHandler: +class EventHandler (ABC): """ Abstract base class for event handler """ __slots__ = () - # this handler wants to know about exceptions before they are reraised by - # the controller - acceptException = False - - def push (self, item): + @abstractmethod + async def push (self, item): raise NotImplementedError () class StatsHandler (EventHandler): __slots__ = ('stats', ) - acceptException = True - def __init__ (self): self.stats = {'requests': 0, 'finished': 0, 'failed': 0, 'bytesRcv': 0} - def push (self, item): - if isinstance (item, Item): + async def push (self, item): + if isinstance (item, RequestResponsePair): self.stats['requests'] += 1 - if item.failed: + if not item.response: self.stats['failed'] += 1 else: self.stats['finished'] += 1 - self.stats['bytesRcv'] += item.encodedDataLength + self.stats['bytesRcv'] += item.response.bytesReceived class LogHandler (EventHandler): """ Handle items by logging information about them """ @@ -83,7 +81,7 @@ class LogHandler (EventHandler): def __init__ (self, logger): self.logger = logger.bind (context=type (self).__name__) - def push (self, item): + 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) @@ -102,6 +100,71 @@ class ControllerStart: 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') + + 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. @@ -110,120 +173,141 @@ class SinglePageController: (stats, warc writer). """ - __slots__ = ('url', 'service', 'behavior', 'settings', 'logger', 'handler') + __slots__ = ('url', 'service', 'behavior', 'settings', 'logger', 'handler', + 'warcinfo', '_enabledBehavior') def __init__ (self, url, logger, \ service, behavior=cbehavior.available, \ - settings=defaultSettings, handler=[]): + settings=defaultSettings, handler=None, \ + warcinfo=None): self.url = url self.service = service self.behavior = behavior self.settings = settings self.logger = logger.bind (context=type (self).__name__, url=url) - self.handler = handler + self.handler = handler or [] + self.warcinfo = warcinfo - def processItem (self, item): + async def processItem (self, item): for h in self.handler: - h.push (item) + await h.push (item) async def run (self): logger = self.logger async def processQueue (): async for item in l: - self.processItem (item) + 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, self.url, logger=logger) as l: + 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)) - start = time.time () + # 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 - enabledBehavior = list (filter (lambda x: self.url in x, + self._enabledBehavior = list (filter (lambda x: self.url in x, map (lambda x: x (l, logger), self.behavior))) - version = await l.tab.Browser.getVersion () + version = await tab.Browser.getVersion () payload = { 'software': getSoftwareInfo (), 'browser': { 'product': version['product'], 'useragent': version['userAgent'], - 'viewport': await getFormattedViewportMetrics (l.tab), + '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'), enabledBehavior)), + 'behavior': list (map (attrgetter('name'), self._enabledBehavior)), + 'insecure': self.settings.insecure, + 'cookies': list (map (lambda x: x.OutputString(), self.settings.cookies)), }, } - self.processItem (ControllerStart (payload)) + if self.warcinfo: + payload['extra'] = self.warcinfo + await self.processItem (ControllerStart (payload)) - await l.start () - for b in enabledBehavior: - async for item in b.onload (): - self.processItem (item) + await l.navigate (self.url) - # wait until the browser has a) been idle for at least - # settings.idleTimeout or b) settings.timeout is exceeded - timeoutProc = asyncio.ensure_future (asyncio.sleep (self.settings.timeout)) - idleTimeout = None + idleProc = asyncio.ensure_future (idle.wait (self.settings.idleTimeout)) while True: - idleProc = asyncio.ensure_future (l.idle.wait ()) try: finished, pending = await asyncio.wait([idleProc, timeoutProc, handle], - return_when=asyncio.FIRST_COMPLETED, timeout=idleTimeout) + return_when=asyncio.FIRST_COMPLETED) except asyncio.CancelledError: idleProc.cancel () timeoutProc.cancel () break - if not finished: - # idle timeout - idleProc.cancel () - timeoutProc.cancel () - break - elif handle in finished: + 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 state change - isIdle = idleProc.result () - if isIdle: - # browser is idle, start the clock - idleTimeout = self.settings.idleTimeout - else: - idleTimeout = None - - for b in enabledBehavior: - async for item in b.onstop (): - self.processItem (item) - await l.tab.Page.stopLoading () + # 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 () - for b in enabledBehavior: - async for item in b.onfinish (): - self.processItem (item) - - # wait until loads from behavior scripts are done - await asyncio.sleep (1) - if not l.idle.get (): - while not await l.idle.wait (): pass + # 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: 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 """ @@ -242,19 +326,17 @@ class DepthLimit (RecursionPolicy): __slots__ = ('maxdepth', ) def __init__ (self, maxdepth=0): - if maxdepth < 0 or maxdepth > 1: - raise ValueError ('Unsupported') 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): """ @@ -271,7 +353,11 @@ class PrefixLimit (RecursionPolicy): 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 hasTemplate (s): + """ Return True if string s has string templates """ + return '{' in s and '}' in s class RecursiveController: """ @@ -281,47 +367,59 @@ class RecursiveController: """ __slots__ = ('url', 'output', 'command', 'logger', 'policy', 'have', - 'pending', 'stats', 'prefix', 'tempdir', 'running', 'concurrency', '_quit') + 'pending', 'stats', 'tempdir', 'running', 'concurrency', + 'copyLock') SCHEME_WHITELIST = {'http', 'https'} - def __init__ (self, url, output, command, logger, prefix='{host}-{date}-', + def __init__ (self, url, output, command, logger, tempdir=None, policy=DepthLimit (0), concurrency=1): self.url = url self.output = output self.command = command - self.prefix = prefix 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} - # initiate graceful shutdown - self._quit = False - async def fetch (self, url): + async def fetch (self, entry, seqnum): """ Fetch a single URL using an external command - command is usually crocoite-grab + command is usually crocoite-single """ + assert isinstance (entry, SetEntry) + + url = entry.value + depth = entry.depth logger = self.logger.bind (url=url) def formatCommand (e): - return e.format (url=url, dest=dest.name) + # provide means to disable variable expansion + if e.startswith ('!'): + return e[1:] + else: + return e.format (url=url, dest=dest.name) - def formatPrefix (p): - return p.format (host=urlparse (url).hostname, date=datetime.utcnow ().isoformat ()) + 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 urlparse (url).scheme not in self.SCHEME_WHITELIST: + if url.scheme not in self.SCHEME_WHITELIST: self.stats['ignored'] += 1 logStats () self.logger.warning ('scheme not whitelisted', url=url, @@ -329,69 +427,115 @@ class RecursiveController: return dest = tempfile.NamedTemporaryFile (dir=self.tempdir, - prefix=formatPrefix (self.prefix), suffix='.warc.gz', + prefix=os.path.basename (self.output) + '-', suffix='.warc.gz', delete=False) - destpath = os.path.join (self.output, os.path.basename (dest.name)) command = list (map (formatCommand, self.command)) - logger.info ('fetch', uuid='1680f384-744c-4b8a-815b-7346e632e8db', command=command, destfile=destpath) - process = await asyncio.create_subprocess_exec (*command, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.DEVNULL, stdin=asyncio.subprocess.DEVNULL, - start_new_session=True) - 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 (removeFragment, 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) + 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 () - code = await process.wait() - if code == 0: - # atomically move once finished - os.rename (dest.name, destpath) - else: - self.stats['crashed'] += 1 - logStats () - - def cancel (self): - """ Gracefully cancel this job, waiting for existing workers to shut down """ - self.logger.info ('cancel', - uuid='d58154c8-ec27-40f2-ab9e-e25c1b21cd88', - pending=len (self.pending), have=len (self.have), - running=len (self.running)) - self._quit = True 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), + pending=len (self.pending), + have=len (self.have)-len(self.running), running=len (self.running)) - self.have = set () - self.pending = set ([self.url]) - - while self.pending and not self._quit: - # 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)) - self.running.add (t) - + 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 () - 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) - - done = asyncio.gather (*self.running) - self.running = set () - log () - |