path: root/python/sandcrawler/ia.py

# XXX: some broken MRO thing going on in here due to python3 object wrangling
# in `wayback` library. Means we can't run pylint.
# pylint: skip-file

import os, sys, time
import requests
import datetime
from collections import namedtuple

import wayback.exception
from http.client import IncompleteRead
from wayback.resourcestore import ResourceStore
from gwb.loader import CDXLoaderFactory

from .misc import b32_hex, requests_retry_session

class CdxApiError(Exception):
    pass

class CdxApiClient:

    def __init__(self, host_url="https://web.archive.org/cdx/search/cdx"):
        self.host_url = host_url
        self.http_session = requests_retry_session(retries=3, backoff_factor=3)
        self.http_session.headers.update({
            'User-Agent': 'Mozilla/5.0 sandcrawler.CdxApiClient',
        })
        self.wayback_endpoint = "https://web.archive.org/web/"

    def lookup_latest(self, url, recent_only=True, follow_redirects=False, redirect_depth=0):
        """
        Looks up most recent HTTP 200 record for the given URL.

        Returns a CDX dict, or None if not found.

        NOTE: could do authorized lookup using cookie to get all fields?
        """

        if redirect_depth >= 15:
            raise CdxApiError("redirect loop (by iteration count)")

        since = datetime.date.today() - datetime.timedelta(weeks=4)
        params = {
            'url': url,
            'matchType': 'exact',
            'limit': -1,
            'output': 'json',
        }
        if recent_only:
            params['from'] = '%04d%02d%02d' % (since.year, since.month, since.day),
        if not follow_redirects:
            params['filter'] = 'statuscode:200'
        resp = self.http_session.get(self.host_url, params=params)
        if resp.status_code != 200:
            raise CdxApiError(resp.text)
        rj = resp.json()
        if len(rj) <= 1:
            return None
        cdx = rj[1]
        assert len(cdx) == 7    # JSON is short
        cdx = dict(
            surt=cdx[0],
            datetime=cdx[1],
            url=cdx[2],
            mimetype=cdx[3],
            http_status=int(cdx[4]),
            sha1b32=cdx[5],
            sha1hex=b32_hex(cdx[5]),
        )
        if follow_redirects and cdx['http_status'] in (301, 302):
            try:
                resp = requests.get(self.wayback_endpoint + cdx['datetime'] + "id_/" + cdx['url'])
            except requests.exceptions.TooManyRedirects:
                raise CdxApiError("redirect loop (wayback fetch)")
            next_url = '/'.join(resp.url.split('/')[5:])
            if next_url == url:
                raise CdxApiError("redirect loop (by url)")
            return self.lookup_latest(next_url, redirect_depth=redirect_depth+1)
        return cdx


class WaybackError(Exception):
    pass

class WaybackClient:

    def __init__(self, cdx_client=None, **kwargs):
        if cdx_client:
            self.cdx_client = cdx_client
        else:
            self.cdx_client = CdxApiClient()
        # /serve/ instead of /download/ doesn't record view count
        self.petabox_base_url = kwargs.get('petabox_base_url', 'http://archive.org/serve/')
        # gwb library will fall back to reading from /opt/.petabox/webdata.secret
        self.petabox_webdata_secret = kwargs.get('petabox_webdata_secret', os.environ.get('PETABOX_WEBDATA_SECRET'))
        self.warc_uri_prefix = kwargs.get('warc_uri_prefix', 'https://archive.org/serve/')
        self.rstore = None

    def fetch_warc_content(self, warc_path, offset, c_size):
        warc_uri = self.warc_uri_prefix + warc_path
        if not self.rstore:
            self.rstore = ResourceStore(loaderfactory=CDXLoaderFactory(
                webdata_secret=self.petabox_webdata_secret,
                download_base_url=self.petabox_base_url))
        try:
            gwb_record = self.rstore.load_resource(warc_uri, offset, c_size)
        except wayback.exception.ResourceUnavailable:
            raise WaybackError("failed to load file contents from wayback/petabox (ResourceUnavailable)")
        except ValueError as ve:
            raise WaybackError("failed to load file contents from wayback/petabox (ValueError: {})".format(ve))
        except EOFError as eofe:
            raise WaybackError("failed to load file contents from wayback/petabox (EOFError: {})".format(eofe))
        except TypeError as te:
            raise WaybackError("failed to load file contents from wayback/petabox (TypeError: {}; likely a bug in wayback python code)".format(te))
        # Note: could consider a generic "except Exception" here, as we get so
        # many petabox errors. Do want jobs to fail loud and clear when the
        # whole cluster is down though.

        if gwb_record.get_status()[0] != 200:
            raise WaybackError("archived HTTP response (WARC) was not 200: {}".format(gwb_record.get_status()[0]))

        try:
            raw_content = gwb_record.open_raw_content().read()
        except IncompleteRead as ire:
            raise WaybackError("failed to read actual file contents from wayback/petabox (IncompleteRead: {})".format(ire))
        return raw_content

    def fetch_warc_by_url_dt(self, url, datetime):
        """
        Helper wrapper that first hits CDX API to get a full CDX row, then
        fetches content from wayback
        """
        cdx_row = self.cdx_client.lookup(url, datetime)
        return self.fetch_warc_content(
            cdx_row['warc_path'],
            cdx_row['warc_offset'],
            cdx_row['warc_csize'])

    def fetch_resource(self, start_url, mimetype=None):
        """
        Looks in wayback for a resource starting at the URL, following any
        redirects. Returns a ResourceResult object.

        In a for loop:

            lookup best CDX
            redirect?
                fetch wayback
                continue
            good?
                fetch wayback
                return success
            bad?
                return failure

        got to end?
            return failure; too many redirects
        """
        next_url = start_url
        urls_seen = [start_url]
        for i in range(25):
            cdx_row = self.cdx_client.lookup_best(next_url, mimetype=mimetype)
            if not cdx_row:
                return None
            if cdx.status_code == 200:
                body = self.fetch_warc_content(cdx.warc_path, cdx.warc_offset, cdx_row.warc_csize)
                return ResourceResult(
                    start_url=start_url,
                    hit=True,
                    status="success",
                    terminal_url=cdx_row.url,
                    terminal_dt=cdx_row.datetime,
                    terminal_status_code=cdx_row.status_code,
                    body=body,
                    cdx=cdx_row,
                )
            elif cdx_row.status_code >= 300 and cdx_row.status_code < 400:
                body = self.fetch_warc_content(cdx_row.warc_path, cdx_row.warc_offset, cdx_row.warc_csize)
                next_url = body.get_redirect_url()
                if next_url in urls_seen:
                    return ResourceResult(
                        start_url=start_url,
                        hit=False,
                        status="redirect-loop",
                        terminal_url=cdx_row.url,
                        terminal_dt=cdx_row.datetime,
                        terminal_status_code=cdx_row.status_code,
                        body=None,
                        cdx=cdx_row,
                    )
                urls_seen.append(next_url)
                continue
            else:
                return ResourceResult(
                    start_url=start_url,
                    hit=False,
                    status="remote-status",
                    terminal_url=cdx_row.url,
                    terminal_dt=cdx_row.datetime,
                    terminal_status_code=cdx_row.status_code,
                    body=None,
                    cdx=cdx_row,
                )
        return ResourceResult(
            start_url=start_url,
            hit=False,
            status="redirects-exceeded",
            terminal_url=cdx_row.url,
            terminal_dt=cdx_row.datetime,
            terminal_status_code=cdx_row.status_code,
            body=None,
            cdx=cdx_row,
        )


class SavePageNowError(Exception):
    pass

SavePageNowResult = namedtuple('SavePageNowResult', [
    'success',
    'status',
    'job_id',
    'request_url',
    'terminal_url',
    'terminal_dt',
    'resources',
])

class SavePageNowClient:

    def __init__(self, v2endpoint="https://web.archive.org/save", **kwargs):
        self.ia_access_key = kwargs.get('ia_access_key',
            os.environ.get('IA_ACCESS_KEY'))
        self.ia_secret_key = kwargs.get('ia_secret_key',
            os.environ.get('IA_SECRET_KEY'))
        self.v2endpoint = v2endpoint
        self.v2_session = requests_retry_session(retries=5, backoff_factor=3)
        self.v2_session.headers.update({
            'User-Agent': 'Mozilla/5.0 sandcrawler.SavePageNowClient',
            'Accept': 'application/json',
            'Authorization': 'LOW {}:{}'.format(self.ia_access_key, self.ia_secret_key),
        })
        self.poll_count = 30
        self.poll_seconds = 3.0

    def save_url_now_v2(self, request_url):
        """
        Returns a "SavePageNowResult" (namedtuple) if SPN request was processed
        at all, or raises an exception if there was an error with SPN itself.

        If SPN2 was unable to fetch the remote content, `success` will be
        false and status will be indicated.

        SavePageNowResult fields:
        - success: boolean if SPN
        - status: "success" or an error message/type
        - job_id: returned by API
        - request_url: url we asked to fetch
        - terminal_url: final primary resource (after any redirects)
        - terminal_timestamp: wayback timestamp of final capture
        - resources: list of all URLs captured

        TODO: parse SPN error codes and handle better. Eg, non-200 remote
        statuses, invalid hosts/URLs, timeouts, backoff, etc.
        """
        if not (self.ia_access_key and self.ia_secret_key):
            raise Exception("SPN2 requires authentication (IA_ACCESS_KEY/IA_SECRET_KEY)")
        resp = self.v2_session.post(
            self.v2endpoint,
            data={
                'url': request_url,
                'capture_all': 1,
                'if_not_archived_within': '1d',
            },
        )
        if resp.status_code != 200:
            raise SavePageNowError("SPN2 status_code: {}, url: {}".format(resp.status_code, request_url))
        resp_json = resp.json()

        if not resp_json or 'job_id' not in resp_json:
            raise SavePageNowError(
                "Didn't get expected 'job_id' field in SPN2 response: {}".format(resp_json))

        job_id = resp_json['job_id']

        # poll until complete
        final_json = None
        for i in range(self.poll_count):
            resp = self.v2_session.get("{}/status/{}".format(self.v2endpoint, resp_json['job_id']))
            try:
                resp.raise_for_status()
            except:
                raise SavePageNowError(resp.content)
            status = resp.json()['status']
            if status == 'pending':
                time.sleep(self.poll_seconds)
            elif status in ('success', 'error'):
                final_json = resp.json()
                break
            else:
                raise SavePageNowError("Unknown SPN2 status:{} url:{}".format(status, request_url))

        if not final_json:
            raise SavePageNowError("SPN2 timed out (polling count exceeded)")

        if final_json['status'] == "success":
            return SavePageNowResult(
                True,
                "success",
                job_id,
                request_url,
                final_json['original_url'],
                final_json['timestamp'],
                final_json['resources'],
            )
        else:
            return SavePageNowResult(
                False,
                final_json.get('status_ext') or final_json['status'],
                job_id,
                request_url,
                None,
                None,
                None,
            )