wrolpi/modules/archive/models.py

import datetime
import json
import pathlib
import re
from typing import Generator, Optional, Iterable

import pytz
from sqlalchemy import Column, Integer, String, ForeignKey, Boolean
from sqlalchemy.orm import relationship, Session, validates
from sqlalchemy.orm.collections import InstrumentedList

from wrolpi.common import ModelHelper, Base, logger
from wrolpi.dates import TZDateTime
from wrolpi.errors import InvalidArchive
from wrolpi.files.lib import split_path_stem_and_suffix
from wrolpi.files.models import File
from wrolpi.media_path import MediaPathType

logger = logger.getChild(__name__)

__all__ = ['Archive', 'Domain']

MATCH_URL = re.compile(r'^\s+?url:\s+?(http.*)', re.MULTILINE)
MATCH_DATE = re.compile(r'^\s+?saved date:\s+?(.*)', re.MULTILINE)


class Archive(Base, ModelHelper):
    __tablename__ = 'archive'
    id = Column(Integer, primary_key=True)

    domain_id = Column(Integer, ForeignKey('domains.id'))
    domain = relationship('Domain', primaryjoin='Archive.domain_id==Domain.id')

    url = Column(String)
    title = Column(String)
    archive_datetime = Column(TZDateTime)
    validated = Column(Boolean)

    # Associated Files.
    singlefile_path: pathlib.Path = Column(MediaPathType, ForeignKey('file.path', ondelete='CASCADE'))
    singlefile_file: File = relationship('File', primaryjoin='Archive.singlefile_path==File.path')
    readability_path: pathlib.Path = Column(MediaPathType, ForeignKey('file.path', ondelete='CASCADE'))
    readability_file: File = relationship('File', primaryjoin='Archive.readability_path==File.path')
    readability_json_path: pathlib.Path = Column(MediaPathType, ForeignKey('file.path', ondelete='CASCADE'))
    readability_json_file: File = relationship('File', primaryjoin='Archive.readability_json_path==File.path')
    readability_txt_path: pathlib.Path = Column(MediaPathType, ForeignKey('file.path', ondelete='CASCADE'))
    readability_txt_file: File = relationship('File', primaryjoin='Archive.readability_txt_path==File.path')
    screenshot_path: pathlib.Path = Column(MediaPathType, ForeignKey('file.path', ondelete='CASCADE'))
    screenshot_file: File = relationship('File', primaryjoin='Archive.screenshot_path==File.path')

    def __repr__(self):
        if self.domain:
            return f'<Archive id={self.id} url={self.url} singlefile={repr(str(self.singlefile_path))} ' \
                   f'domain={self.domain.domain}>'
        return f'<Archive id={self.id} url={self.url} singlefile={repr(str(self.singlefile_path))}>'

    def my_paths(self) -> Generator[pathlib.Path, None, None]:
        if self.singlefile_path:
            yield self.singlefile_path
        if self.readability_path:
            yield self.readability_path
        if self.readability_json_path:
            yield self.readability_json_path
        if self.readability_txt_path:
            yield self.readability_txt_path
        if self.screenshot_path:
            yield self.screenshot_path

    def my_files(self) -> Generator[File, None, None]:
        if self.singlefile_file:
            yield self.singlefile_file
        if self.readability_file:
            yield self.readability_file
        if self.readability_json_file:
            yield self.readability_json_file
        if self.readability_txt_file:
            yield self.readability_txt_file
        if self.screenshot_file:
            yield self.screenshot_file

    def unlink(self):
        """
        Remove any files in this archive.
        """
        # Don't fail if a file is missing, it could have been deleted manually.
        for path in self.my_paths():
            path.unlink(missing_ok=True)

    def __gt__(self, other) -> bool:
        """
        Compare Archives according to their archive datetime.
        """
        if not isinstance(other, Archive):
            raise ValueError(f'Cannot compare {type(other)} to Archive!')

        return self.archive_datetime > other.archive_datetime

    def __json__(self):
        stem, _ = split_path_stem_and_suffix(self.singlefile_path) if self.singlefile_path else (None, None)
        d = dict(
            archive_datetime=self.archive_datetime,
            domain=self.domain.dict() if self.domain else None,
            domain_id=self.domain_id,
            id=self.id,
            readability_json_path=self.readability_json_path,
            readability_path=self.readability_path,
            readability_txt_path=self.readability_txt_path,
            screenshot_path=self.screenshot_path,
            singlefile_path=self.singlefile_path,
            stem=stem,
            title=self.title,
            url=self.url,
        )
        return d

    def delete(self):
        self.unlink()

        session = Session.object_session(self)

        for file in self.my_files():
            session.delete(file)

        session.delete(self)

        if self.domain:
            domain_archives = [i.id for i in self.domain.archives]
            if not domain_archives:
                self.domain.delete()

    @property
    def history(self) -> Iterable[Base]:
        """Get a list of Archives that share my URL."""
        session = Session.object_session(self)
        history = list(session.query(Archive).filter(
            Archive.id != self.id,
            Archive.url == self.url,
            Archive.url != None,  # noqa
        ).order_by(Archive.archive_datetime))
        return history

    def read_singlefile_data(self):
        """
        Read the start of the singlefile (if any) and decode the Archive information.
        """
        path = self.singlefile_file.path if self.singlefile_file else self.singlefile_path
        if not path:
            # Can't read contents of nothing.
            raise ValueError('Cannot read singlefile data when this has no files')

        with path.open('rt') as fh:
            head = fh.read(1000)
            if 'Page saved with SingleFile' not in head:
                return
            try:
                if match := MATCH_URL.findall(head):
                    self.url = match[0].strip()
            except Exception as e:
                logger.error(f'Could not get URL from singlefile {path}', exc_info=e)
            try:
                if match := MATCH_DATE.findall(head):
                    dt = match[0].strip()
                    dt = ' '.join(dt.split(' ')[:5])
                    # SingleFile uses GMT.
                    dt = datetime.datetime.strptime(
                        dt,
                        '%a %b %d %Y %H:%M:%S'  # Fri Jun 17 2022 19:24:52
                    ).replace(tzinfo=pytz.timezone('GMT'))
                    self.archive_datetime = dt
            except Exception as e:
                logger.error(f'Could not get archive date from singlefile {path}', exc_info=e)

    def read_readability_data(self):
        """Read the Readability JSON file, apply its contents to this record."""
        readability_json_file = self.readability_json_file
        if not readability_json_file or not readability_json_file.path.is_file():
            path = self.singlefile_file.path if self.singlefile_file else self.singlefile_path
            logger.warning(f'{path} does not have an info json file')
            return

        try:
            with readability_json_file.path.open() as fh:
                json_contents = json.load(fh)
                url = json_contents.get('url')
                title = json_contents.get('title')
        except Exception as e:
            raise InvalidArchive() from e

        self.url = url
        self.title = title

    def apply_domain(self):
        """Get the domain from the URL."""
        from modules.archive.lib import get_or_create_domain
        domain = None
        if self.url:
            domain = get_or_create_domain(Session.object_session(self), self.url)
        # Clear domain if the URL is missing.
        self.domain_id = domain.id if domain else None

    def apply_title(self):
        """Get the title from the Singlefile, if its missing."""
        from modules.archive.lib import get_title_from_html
        if self.singlefile_path and not self.title:
            self.title = get_title_from_html(self.singlefile_path.read_text())

    def validate(self):
        try:
            self.read_readability_data()
            self.read_singlefile_data()
            self.apply_domain()
            self.apply_title()
            self.validated = True
        except Exception as e:
            logger.warning(f'Unable to validate {self}', exc_info=e)

    @staticmethod
    def find_by_path(path, session: Session) -> Optional[Base]:
        archive = session.query(Archive).filter_by(singlefile_path=path).one_or_none()
        return archive

    @staticmethod
    def find_by_paths(paths, session):
        archives = list(session.query(Archive).filter(Archive.singlefile_path.in_(paths)))
        return archives

    @property
    def primary_path(self):
        return self.singlefile_path


class Domain(Base, ModelHelper):
    __tablename__ = 'domains'  # plural to avoid conflict
    id = Column(Integer, primary_key=True)

    domain = Column(String)
    directory = Column(MediaPathType)

    archives: InstrumentedList = relationship('Archive', primaryjoin='Archive.domain_id==Domain.id')

    def __repr__(self):
        return f'<Domain id={self.id} domain={self.domain} directory={self.directory}>'

    def delete(self):
        session = Session.object_session(self)
        session.execute('DELETE FROM archive WHERE domain_id=:id', dict(id=self.id))
        session.query(Domain).filter_by(id=self.id).delete()

    @validates('domain')
    def validate_domain(self, key, value: str):
        if not isinstance(value, str):
            raise ValueError('Domain must be a string')
        if len(value.split('.')) < 2:
            raise ValueError(f'Domain must contain at least one "." domain={repr(value)}')
        return value