Finalize proposal of cert download and convert

The proposed method of processing certificates is breadth-first.
By that I mean that single task is done for all certificates, then next
task is evaluated. All tasks should be implemented as static methods
and are to be paralelized by calling `sec-certs/cert_processing.py`.
Currently, a generic method `process_paralel()` sits there that is
capable of pretty much arbitrary parallel processing (both using threads
and processes). Any methods to-be called by this method must require
unpacked arguments on the input.

5 files changed, 97 insertions, 51 deletions

diff --git a/cc_oop_demo.py b/cc_oop_demo.py
index f8ae5508..fae5d911 100644
--- a/cc_oop_demo.py
+++ b/cc_oop_demo.py
@@ -27,6 +27,12 @@ def main():
 
     assert dset == new_dset
 
+    # Download pdfs
+    dset.download_all_pdfs()
+
+    # Convert pdfs to text
+    dset.convert_all_pdfs()
+
     end = datetime.now()
     logging.info(f'The computation took {(end-start)} seconds.')
 
diff --git a/sec_certs/cert_processing.py b/sec_certs/cert_processing.py
new file mode 100644
index 00000000..6d0f4d92
--- /dev/null
+++ b/sec_certs/cert_processing.py
@@ -0,0 +1,28 @@
+from tqdm import tqdm
+from multiprocessing.pool import Pool, ThreadPool
+import sec_certs.constants as constants
+import logging
+import time
+
+# TODO: Add timeout. Kinda meh with ThreadingTimeout, SignalTimeout does not work on Windows, stopit package.
+def process_parallel(func, items, max_workers, callback=None, use_threading=True, progress_bar=True):
+    if use_threading is True:
+        pool = ThreadPool(max_workers)
+    else:
+        pool = Pool(max_workers)
+
+    results = [pool.apply_async(func, (*i, ), callback=callback) for i in items]
+
+    if progress_bar is True:
+        bar = tqdm(total=len(results))
+        while not all([x.ready() for x in results]):
+            done_count = len([x.ready() for x in results if x.ready()])
+            bar.update(done_count - bar.n)
+            time.sleep(1)
+        bar.update(len(results) - bar.n)
+        bar.close()
+
+    pool.close()
+    pool.join()
+
+    return [r.get() for r in results]
\ No newline at end of file
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index efa2740b..703aab88 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -1,6 +1,6 @@
 import os
 import re
-from datetime import datetime, time
+from datetime import datetime
 import locale
 import logging
 from typing import Dict, List, ClassVar, Collection
@@ -11,9 +11,9 @@ import copy
 from abc import ABC, abstractmethod
 from pathlib import Path
 import shutil
-from multiprocessing import Pool, pool
-import tqdm
+
 from functools import partial
+import requests
 
 
 from tabula import read_pdf
@@ -28,6 +28,7 @@ from sec_certs.extract_certificates import extract_certificates_keywords
 from sec_certs.constants import FIPS_NOT_AVAILABLE_CERT_SIZE
 import sec_certs.constants as constants
 import sec_certs.download as download
+import sec_certs.cert_processing as cert_processing
 
 
 class Dataset(ABC):
@@ -73,7 +74,7 @@ class Dataset(ABC):
     @classmethod
     def from_dict(cls, dct: Dict):
         certs = {x.dgst: x for x in dct['certs']}
-        dset =  cls(certs, Path(dct['root_dir']), dct['name'], dct['description'])
+        dset = cls(certs, Path(dct['root_dir']), dct['name'], dct['description'])
         assert len(dset) == dct['n_certs']
         return dset
 
@@ -106,24 +107,41 @@ class Dataset(ABC):
             f'Added {len(will_be_added)} new and merged further {n_merged} certificates to the dataset.')
 
     @staticmethod
-    def convert_pdfs_to_text(pdf_paths: Collection[Path], txt_paths: Collection[Path]):
+    def convert_pdfs_to_txt(pdf_paths: Collection[Path], txt_paths: Collection[Path]):
         assert len(pdf_paths) == len(txt_paths)
-        results = []
+
         partial_convert_pdf = partial(helpers.convert_pdf_file, options=['-raw'])
-        with tqdm.tqdm(total=len(pdf_paths)) as progress:
-            for result in pool.ThreadPool(constants.N_THREADS).imap(partial_convert_pdf, zip(pdf_paths, txt_paths)):
-                progress.update(1)
-                results.append(result)
+        exit_codes = cert_processing.process_parallel(partial_convert_pdf,
+                                                      list(zip(pdf_paths, txt_paths)),
+                                                      constants.N_THREADS)
+
+        n_successful = len([e for e in exit_codes if e == constants.RETURNCODE_OK])
+        logging.info(f'Successfully converted {n_successful} files pdf->txt, {len(exit_codes) - n_successful} failed.')
+
+        for path, e in zip(pdf_paths, exit_codes):
+            if e != constants.RETURNCODE_OK:
+                logging.info(f'Failed to convert {path}, exit code: {e}')
 
     @staticmethod
-    def get_corrupted_pdfs(pdf_paths):
-        return [p for p in pdf_paths if p.stat().st_size < constants.MIN_CORRECT_CERT_SIZE]
+    def download_parallel(urls, paths, prune_corrupted=True):
+        exit_codes = cert_processing.process_parallel(download.download_file,
+                                                      list(zip(urls, paths)),
+                                                      constants.N_THREADS)
+        n_successful = len([e for e in exit_codes if e == requests.codes.ok])
+        logging.info(f'Successfully downloaded {n_successful} files, {len(exit_codes) - n_successful} failed.')
 
+        for url, e in zip(urls, exit_codes):
+            if e != requests.codes.ok:
+                logging.error(f'Failed to download {url}, exit code: {e}')
+
+        if prune_corrupted is True:
+            for p in paths:
+                if p.exists() and p.stat().st_size < constants.MIN_CORRECT_CERT_SIZE:
+                    logging.error(f'Corrupted file at: {p}')
+                    # TODO: Delete
 
-class CCDataset(Dataset):
-    def __init__(self, certs, root_dir, name, description):
-        super().__init__(certs, root_dir, name, description)
 
+class CCDataset(Dataset):
     @property
     def web_dir(self) -> Path:
         return self.root_dir / 'web'
@@ -210,10 +228,13 @@ class CCDataset(Dataset):
             html_items = [x for x in html_items if 'archived' not in str(x[1])]
             csv_items = [x for x in csv_items if 'archived' not in str(x[1])]
 
+        html_urls, html_paths = [x[0] for x in html_items], [x[1] for x in html_items]
+        csv_urls, csv_paths = [x[0] for x in csv_items], [x[1] for x in csv_items]
+
         if to_download is True:
             logging.info('Downloading required csv and html files.')
-            helpers.download_parallel(html_items, num_threads=8)
-            helpers.download_parallel(csv_items, num_threads=8)
+            self.download_parallel(html_urls, html_paths)
+            self.download_parallel(csv_urls, csv_paths)
 
         logging.info('Adding CSV certificates to CommonCriteria dataset.')
         csv_certs = self.get_all_certs_from_csv(get_active, get_archived)
@@ -387,39 +408,36 @@ class CCDataset(Dataset):
     def download_reports(self):
         self.reports_pdf_dir.mkdir(parents=True, exist_ok=True)
         reports_urls = [x.report_link for x in self]
-        download.download_parallel(list(zip(reports_urls, self.report_pdf_paths.values())), constants.N_THREADS)
+        self.download_parallel(reports_urls, self.report_pdf_paths.values(), prune_corrupted=True)
 
     def download_targets(self):
         self.targets_pdf_dir.mkdir(parents=True, exist_ok=True)
         target_urls = [x.st_link for x in self]
-        download.download_parallel(list(zip(target_urls, self.target_pdf_paths.values())), constants.N_THREADS)
+        self.download_parallel(target_urls, self.target_pdf_paths.values(), prune_corrupted=True)
 
     def download_all_pdfs(self):
         logging.info('Downloading CC certificate reports')
         self.download_reports()
 
-        # TODO: Do checks below live when downloading and re-download straight away?
-        corrupted_reports = self.get_corrupted_pdfs(self.report_pdf_paths.values())
-        for r in corrupted_reports:
-            logging.error(f'Corrupted pdf file at: {r}')
-
         logging.info('Downloading CC security targets')
         self.download_targets()
 
-        # TODO: Do checks below live when downloading and re-download straight away?
-        corrupted_targets = self.get_corrupted_pdfs(self.target_pdf_paths.values())
-        for t in corrupted_targets:
-            logging.error(f'Corrupted pdf file at: {t}')
+    def convert_reports_to_txt(self):
+        self.reports_txt_dir.mkdir(parents=True, exist_ok=True)
+        # TODO: Get rid of the list() invocation here.
+        self.convert_pdfs_to_txt(list(self.report_pdf_paths.values()), list(self.report_txt_paths.values()))
 
-    def convert_all_pdfs(self):
+    def convert_targets_to_txt(self):
+        self.targets_txt_dir.mkdir(parents=True, exist_ok=True)
         # TODO: Get rid of the list() invocation here.
+        self.convert_pdfs_to_txt(list(self.target_pdf_paths.values()), list(self.target_txt_paths.values()))
+
+    def convert_all_pdfs(self):
         logging.info('Converting CC certificate reports to .txt')
-        self.reports_txt_dir.mkdir(parents=True, exist_ok=True)
-        self.convert_pdfs_to_text(list(self.report_pdf_paths.values()), list(self.report_txt_paths.values()))
+        self.convert_reports_to_txt()
 
         logging.info('Converting CC security targets to .txt')
-        self.targets_txt_dir.mkdir(parents=True, exist_ok=True)
-        self.convert_pdfs_to_text(list(self.target_pdf_paths.values()), list(self.target_txt_paths.values()))
+        self.convert_targets_to_txt()
 
 
 class FIPSDataset(Dataset):
diff --git a/sec_certs/download.py b/sec_certs/download.py
index d01da79f..fc225150 100644
--- a/sec_certs/download.py
+++ b/sec_certs/download.py
@@ -1,5 +1,4 @@
 import os
-import logging
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
 from tqdm import tqdm
@@ -7,8 +6,7 @@ from typing import Sequence, Tuple, List
 
 import requests
 
-from .files import search_files
-import sec_certs.constants as constants
+from sec_certs.files import search_files
 
 CC_WEB_URL = 'https://www.commoncriteriaportal.org'
 
@@ -31,8 +29,6 @@ def download_parallel(items: Sequence[Tuple[str, Path]], num_threads: int) -> Se
         for response in pool.imap(download, items):
             progress.update(1)
             responses.append(response)
-            if response[1] != constants.RESPONSE_OK:
-                logging.error(f'Request for url {response[0]} returned {response[1]}')
     pool.close()
     pool.join()
     return responses
diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
index 718b66cd..942a37fd 100644
--- a/sec_certs/helpers.py
+++ b/sec_certs/helpers.py
@@ -13,16 +13,17 @@ from datetime import date
 import numpy as np
 import pandas as pd
 import subprocess
-from bs4 import Tag, NavigableString
-
-import sec_certs.constants as constants
 
 
 def download_file(url: str, output: Path) -> int:
-    r = requests.get(url, allow_redirects=True)
-    with output.open('wb') as f:
-        f.write(r.content)
-    return r.status_code
+    try:
+        r = requests.get(url, allow_redirects=True, timeout=5)
+        if r.status_code == requests.codes.ok:
+            with output.open("wb") as f:
+                f.write(r.content)
+        return r.status_code
+    except requests.exceptions.Timeout:
+        return requests.codes.timeout
 
 
 def download_parallel(items: Sequence[Tuple[str, Path]], num_threads: int) -> Sequence[Tuple[str, int]]:
@@ -158,12 +159,9 @@ def repair_pdf(file: Path):
     pdf.save(file)
 
 
-def convert_pdf_file(filepaths: Tuple[Path, Path], options):
-    pdf_path, txt_path = filepaths[0], filepaths[1]
-    proc_result = subprocess.run(['pdftotext', *options, pdf_path, txt_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-    if proc_result.returncode != constants.RETURNCODE_OK:
-        logging.error(f'Converting pdf {pdf_path} resulted into the following result: {proc_result}')
-    return proc_result
+def convert_pdf_file(pdf_path: Path, txt_path: Path, options):
+    return subprocess.run(['pdftotext', *options, pdf_path, txt_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode
+