nistsp800.py

#!/usr/bin/env python3.8


import sqlite3
import itertools
import threading
import webbrowser

from pathlib import Path
from functools import partial
from urllib.parse import urlparse
from contextlib import contextmanager
from concurrent.futures import ThreadPoolExecutor

import tkinter as tk

import requests
from PyPDF2 import PdfFileReader
from bs4 import BeautifulSoup

PDF_DIR = Path('NISTSP800')
PDF_DB = PDF_DIR / 'NISTSP800.db'
if not PDF_DIR.exists():
    PDF_DIR.mkdir()
NIST_SP800_HOMEPAGE = 'https://csrc.nist.gov/publications/sp800'



def scrape_nist():
    
    with getdb() as db:
        c = db.cursor()
        with requests.Session() as sess:
            (r:=sess.get(NIST_SP800_HOMEPAGE)).raise_for_status()
            homepage = BeautifulSoup(r.content, 'html.parser')
            for row in homepage.select('table#publications-results-table tbody tr'):
                title = row.select('td div a')[0].text
                download = row.select('td span a[href^="https://doi"]')
                download = download or row.select('td span a[href$=".pdf"]')
                if not download:
                    print('no links:', title)
                    continue
                print('processing', title, '. . .', end=' ')
                url = download[0]['href']
                
                while True:
                    try:
                        file, url = getpdf(url)
                        c.execute(
                            'update pdfs set url=?, title=? where file=?',
                            (url, title, file)
                        )
                        db.commit()
                        insertpdf(db, file, url)
                        break
                    except FileExistsError:
                        print('SKIPPED (file exists)')
                        continue
                    
                    except Exception as e:
                        print('\nERROR:', e)
                        try:
                            i = input('try again or skip').lower()
                        except (KeyboardInterrupt, EOFError):
                            print('STOPPED')
                            return
                        if 't' in i:
                            continue
                        if 's' in i:
                            print('SKIPPED (Manual)')
                            break
                        raise e
                    
                print('OK')


def dereference_url(fragment):
    if not isinstance(fragment, str):
        fragment = fragment.geturl()
    url = urlparse(NIST_SP800_HOMEPAGE)
    url = url._replace(path=fragment)
    return url


def getpdf(url):
    url = urlparse(url)  # in case of redirect
    if not url.netloc:
        url = dereference_url(url)
    # only get headers until file name has been confirmed
    (r:=requests.get(url.geturl(), stream=True)).raise_for_status()
    url = urlparse(r.url)  # in case of redirect
    file_name = url.path.split('/')[-1]

    # TODO: delete this line
    return str(PDF_DIR / file_name), url.geturl()

    with (PDF_DIR / file_name).open('xb') as f:
        assert file_name.lower().endswith('pdf')
        f.write(r.content)
    
    return str(PDF_DIR / file_name), url.geturl()


def getdb():
    db = sqlite3.connect(
        str(PDF_DB), 
        detect_types=sqlite3.PARSE_DECLTYPES,
        check_same_thread=False
    )
    db.row_factory = sqlite3.Row
    db.executescript('''
        BEGIN;
        CREATE TABLE IF NOT EXISTS pdfs (
            file TEXT NOT NULL,
            url TEXT NOT NULL
        );
        
        CREATE TABLE IF NOT EXISTS pages (
            number INTEGER NOT NULL,
            content TEXT NOT NULL,
            
            pdf_id INTEGER NOT NULL,
            FOREIGN KEY (pdf_id)
                REFERENCES pdfs (id) 
                ON DELETE CASCADE
        );
        
        COMMIT;
    ''')
    return db

def insertpdf(db, title, file, url):    
    with db:
        db.execute('''
            INSERT INTO pdfs(file, url) 
            SELECT ?, ?, ?
            WHERE NOT EXISTS(
                SELECT 1 FROM pdfs WHERE title=? and file = ? and url = ?
            );
            ''', (title, file, url)*2
        )
    
    
        pdf_id = db.execute(
            'SELECT rowid FROM pdfs WHERE file = ?', 
            (file,)
        ).fetchone()['rowid']
        
        with open(file, 'rb') as f:
            pdf = PdfFileReader(f)
            pages = zip(
                itertools.count(1),
                (page.extractText() for page in pdf.pages),
                itertools.cycle([pdf_id])
            )
            db.executemany(
                'INSERT INTO pages(number, content, pdf_id) VALUES (?,?,?)', 
                pages
            )
        
@contextmanager
def search_context():
    def _query_wrapper(c, query):
        results = c.execute('''
            SELECT DISTINCT page.number, pdf.file, pdf.title
            FROM all_pages AS page 
            JOIN pdfs AS pdf ON pdf.rowid == page.pdf_id
            WHERE page.content MATCH ?
            ORDER BY rank
            ''',
            (query,)
        ).fetchall()
        return [(r['number'], r['file'], r['title']) for r in results]
    
    db = getdb()
    with db:
        db.executescript('''
            DROP TABLE IF EXISTS all_pages;
            CREATE VIRTUAL TABLE all_pages
                USING FTS5(number, content, pdf_id);
            
            INSERT INTO all_pages SELECT * FROM pages;
        ''')
        try:
            yield partial(_query_wrapper, db.cursor())
        finally:
            db.execute('DROP TABLE all_pages')
            db.commit()
        


class App(tk.Frame):
    def __init__(self, master=None, **kwargs):
        self.search = kwargs.pop('search')
        super().__init__(master, **kwargs)
        self.lock = threading.Lock()
        self.executor = ThreadPoolExecutor()
        self._results = None
        
        self.searchbox = tk.Entry(self)
        self.searchbox.pack(pady=10)
        self.searchbox.bind('<Key>', self.start_search)
        
        self.listbox = tk.Listbox(self)
        self.listbox.pack(fill=tk.BOTH, expand=tk.YES)
        self.listbox.bind('<Double-Button-1>', self.open_link)


    def start_search(self, _event=None):
        self.after(1, self._start_search)  # dont ask

    def _start_search(self):
        if self.lock.locked():
            return
        
        query = self.searchbox.get()
        if query:
            self.lock.acquire()
            fut = self.executor.submit(self.search, query)
            fut.add_done_callback(self._finish_search)
        
        
    def _finish_search(self, fut):
        self._results = fut.result()
        self.listbox.delete(1, self.listbox.size()+1)
        for i, page in enumerate(self._results, 1):
            self.listbox.insert(i, f'{page[0]} {page[2]}')
        self.lock.release()

    
    def open_link(self, _event=None):
        selection = self.listbox.curselection()
        if selection:
            breakpoint()
            selection = selection[0]
            string = self.listbox.get(selection)
            page, title = string.split(' ', 0)
            file = ''
            if self._results:
                for res in self._results:
                    if title in res:
                        file = res[1]
                        file = str(Path(file).absolute())
                        url = f'file://{file}#page={page}'            
                        webbrowser.open(url)
                        return


def main():
    with search_context() as search:
        app = App(search=search)
        app.pack(padx=10, pady=10)
        app.mainloop()
    

if __name__ == '__main__':
    main()