NDOJ/judge/pdf_problems.py

import base64
import errno
import io
import json
import logging
import os
import shutil
import subprocess
import uuid

from django.conf import settings
from django.utils.translation import gettext

logger = logging.getLogger('judge.problem.pdf')

HAS_SELENIUM = False
if settings.USE_SELENIUM:
    try:
        from selenium import webdriver
        from selenium.common.exceptions import TimeoutException
        from selenium.webdriver.common.by import By
        from selenium.webdriver.support import expected_conditions as EC
        from selenium.webdriver.support.ui import WebDriverWait
        HAS_SELENIUM = True
    except ImportError:
        logger.warning('Failed to import Selenium', exc_info=True)

HAS_PHANTOMJS = os.access(settings.PHANTOMJS, os.X_OK)
HAS_SLIMERJS = os.access(settings.SLIMERJS, os.X_OK)

NODE_PATH = settings.NODEJS
PUPPETEER_MODULE = settings.PUPPETEER_MODULE
HAS_PUPPETEER = os.access(NODE_PATH, os.X_OK) and os.path.isdir(PUPPETEER_MODULE)

HAS_PDF = (os.path.isdir(settings.DMOJ_PDF_PROBLEM_CACHE) and
           (HAS_PHANTOMJS or HAS_SLIMERJS or HAS_PUPPETEER or HAS_SELENIUM))

EXIFTOOL = settings.EXIFTOOL
HAS_EXIFTOOL = os.access(EXIFTOOL, os.X_OK)


class BasePdfMaker(object):
    math_engine = 'jax'
    title = None

    def __init__(self, dir=None, clean_up=True):
        self.dir = dir or os.path.join(settings.DMOJ_PDF_PROBLEM_TEMP_DIR, str(uuid.uuid1()))
        self.proc = None
        self.log = None
        self.htmlfile = os.path.join(self.dir, 'input.html')
        self.pdffile = os.path.join(self.dir, 'output.pdf')
        self.clean_up = clean_up

    def load(self, file, source):
        with open(os.path.join(self.dir, file), 'w') as target, open(source) as source:
            target.write(source.read())

    def make(self, debug=False):
        self._make(debug)

        if self.title and HAS_EXIFTOOL:
            try:
                subprocess.check_output([EXIFTOOL, '-Title=%s' % (self.title,), self.pdffile])
            except subprocess.CalledProcessError as e:
                logger.error('Failed to run exiftool to set title for: %s\n%s', self.title, e.output)

    def _make(self, debug):
        raise NotImplementedError()

    @property
    def html(self):
        with io.open(self.htmlfile, encoding='utf-8') as f:
            return f.read()

    @html.setter
    def html(self, data):
        with io.open(self.htmlfile, 'w', encoding='utf-8') as f:
            f.write(data)

    @property
    def success(self):
        return self.proc.returncode == 0

    @property
    def created(self):
        return os.path.exists(self.pdffile)

    def __enter__(self):
        try:
            os.makedirs(self.dir)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.clean_up:
            shutil.rmtree(self.dir, ignore_errors=True)


class PhantomJSPdfMaker(BasePdfMaker):
    template = '''\
"use strict";
var page = require('webpage').create();
var param = {params};

page.paperSize = {
    format: param.paper, orientation: 'portrait', margin: '1cm',
    footer: {
        height: '1cm',
        contents: phantom.callback(function(num, pages) {
            return ('<center style="margin: 0 auto; font-family: Segoe UI; font-size: 10px">'
                  + param.footer.replace('[page]', num).replace('[topage]', pages) + '</center>');
        })
    }
};

page.onCallback = function (data) {
    if (data.action === 'snapshot') {
        page.render(param.output);
        phantom.exit();
    }
}

page.open(param.input, function (status) {
    if (status !== 'success') {
        console.log('Unable to load the address!');
        phantom.exit(1);
    } else {
        page.evaluate(function (zoom) {
            document.documentElement.style.zoom = zoom;
        }, param.zoom);
        window.setTimeout(function () {
            page.render(param.output);
            phantom.exit();
        }, param.timeout);
    }
});
'''

    def get_render_script(self):
        return self.template.replace('{params}', json.dumps({
            'zoom': settings.PHANTOMJS_PDF_ZOOM,
            'timeout': int(settings.PHANTOMJS_PDF_TIMEOUT * 1000),
            'input': 'input.html', 'output': 'output.pdf',
            'paper': settings.PHANTOMJS_PAPER_SIZE,
            'footer': gettext('Page [page] of [topage]'),
        }))

    def _make(self, debug):
        with io.open(os.path.join(self.dir, '_render.js'), 'w', encoding='utf-8') as f:
            f.write(self.get_render_script())
        cmdline = [settings.PHANTOMJS, '_render.js']
        self.proc = subprocess.Popen(cmdline, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=self.dir)
        self.log = self.proc.communicate()[0]


class SlimerJSPdfMaker(BasePdfMaker):
    math_engine = 'mml'

    template = '''\
"use strict";
try {
    var param = {params};

    var {Cc, Ci} = require('chrome');
    var prefs = Cc['@mozilla.org/preferences-service;1'].getService(Ci.nsIPrefService);
    // Changing the serif font so that printed footers show up as Segoe UI.
    var branch = prefs.getBranch('font.name.serif.');
    branch.setCharPref('x-western', 'Segoe UI');

    var page = require('webpage').create();

    page.paperSize = {
        format: param.paper, orientation: 'portrait', margin: '1cm', edge: '0.5cm',
        footerStr: { left: '', right: '', center: param.footer }
    };

    page.open(param.input, function (status) {
        if (status !== 'success') {
            console.log('Unable to load the address!');
            slimer.exit(1);
        } else {
            page.render(param.output, { ratio: param.zoom });
            slimer.exit();
        }
    });
} catch (e) {
    console.error(e);
    slimer.exit(1);
}
'''

    def get_render_script(self):
        return self.template.replace('{params}', json.dumps({
            'zoom': settings.SLIMERJS_PDF_ZOOM,
            'input': 'input.html', 'output': 'output.pdf',
            'paper': settings.SLIMERJS_PAPER_SIZE,
            'footer': gettext('Page [page] of [topage]').replace('[page]', '&P').replace('[topage]', '&L'),
        }))

    def _make(self, debug):
        with io.open(os.path.join(self.dir, '_render.js'), 'w', encoding='utf-8') as f:
            f.write(self.get_render_script())

        env = None
        firefox = settings.SLIMERJS_FIREFOX_PATH
        if firefox:
            env = os.environ.copy()
            env['SLIMERJSLAUNCHER'] = firefox

        cmdline = [settings.SLIMERJS, '--headless', '_render.js']
        self.proc = subprocess.Popen(cmdline, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=self.dir, env=env)
        self.log = self.proc.communicate()[0]


class PuppeteerPDFRender(BasePdfMaker):
    template = '''\
"use strict";
const param = {params};
const puppeteer = require('puppeteer');

puppeteer.launch().then(browser => Promise.resolve()
    .then(async () => {
        const page = await browser.newPage();
        await page.goto(param.input, { waitUntil: 'load' });
        await page.waitForSelector('.math-loaded', { timeout: 15000 });
        await page.pdf({
            path: param.output,
            format: param.paper,
            margin: {
                top: '1cm',
                bottom: '1cm',
                left: '1cm',
                right: '1cm',
            },
            printBackground: true,
            displayHeaderFooter: true,
            headerTemplate: '<div></div>',
            footerTemplate: '<center style="margin: 0 auto; font-family: Segoe UI; font-size: 10px">' +
                param.footer.replace('[page]', '<span class="pageNumber"></span>')
                            .replace('[topage]', '<span class="totalPages"></span>')
                + '</center>',
        });
        await browser.close();
    })
    .catch(e => browser.close().then(() => {throw e}))
).catch(e => {
    console.error(e);
    process.exit(1);
});
'''

    def get_render_script(self):
        return self.template.replace('{params}', json.dumps({
            'input': 'file://%s' % self.htmlfile,
            'output': self.pdffile,
            'paper': settings.PUPPETEER_PAPER_SIZE,
            'footer': gettext('Page [page] of [topage]'),
        }))

    def _make(self, debug):
        with io.open(os.path.join(self.dir, '_render.js'), 'w', encoding='utf-8') as f:
            f.write(self.get_render_script())

        env = os.environ.copy()
        env['NODE_PATH'] = os.path.dirname(PUPPETEER_MODULE)

        cmdline = [NODE_PATH, '_render.js']
        self.proc = subprocess.Popen(cmdline, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=self.dir, env=env)
        self.log = self.proc.communicate()[0]

class SeleniumPDFRender(BasePdfMaker):
    success = False
    template = {
        'printBackground': True,
        'displayHeaderFooter': True,
        'headerTemplate': '<div></div>',
        'footerTemplate': '<center style="margin: 0 auto; font-family: Segoe UI; font-size: 10px">' +
                          gettext('Page %s of %s') %
                          ('<span class="pageNumber"></span>', '<span class="totalPages"></span>') +
                          '</center>',
    }

    def get_log(self, driver):
        return '\n'.join(map(str, driver.get_log('driver') + driver.get_log('browser')))

    def _make(self, debug):
        options = webdriver.ChromeOptions()
        options.add_argument("--headless")
        options.add_argument("--no-sandbox") # for root
        options.binary_location = settings.SELENIUM_CUSTOM_CHROME_PATH

        browser = webdriver.Chrome(settings.SELENIUM_CHROMEDRIVER_PATH, options=options)
        browser.get('file://%s' % self.htmlfile)
        self.log = self.get_log(browser)

        try:
            WebDriverWait(browser, 15).until(EC.presence_of_element_located((By.CLASS_NAME, 'math-loaded')))
        except TimeoutException:
            logger.error('PDF math rendering timed out')
            self.log = self.get_log(browser) + '\nPDF math rendering timed out'
            browser.quit()
            return
        response = browser.execute_cdp_cmd('Page.printToPDF', self.template)
        self.log = self.get_log(browser)
        if not response:
            browser.quit()
            return

        with open(self.pdffile, 'wb') as f:
            f.write(base64.b64decode(response['data']))

        self.success = True
        browser.quit()


if HAS_PUPPETEER:
    DefaultPdfMaker = PuppeteerPDFRender
elif HAS_SELENIUM:
    DefaultPdfMaker = SeleniumPDFRender
elif HAS_SLIMERJS:
    DefaultPdfMaker = SlimerJSPdfMaker
elif HAS_PHANTOMJS:
    DefaultPdfMaker = PhantomJSPdfMaker
else:
    DefaultPdfMaker = None