diff --git a/pywb/default_config.yaml b/pywb/default_config.yaml index 2653a96c..c031928a 100644 --- a/pywb/default_config.yaml +++ b/pywb/default_config.yaml @@ -14,12 +14,19 @@ paths: template_files: banner_html: banner.html - head_insert: head_insert.html - frame_insert: frame_insert.html + head_insert_html: head_insert.html + frame_insert_html: frame_insert.html query_html: query.html search_html: search.html + not_found_html: not_found.html + shared_template_files: + home_html: index.html + error_html: error.html + + proxy_cert_download_html: proxy_cert_download.html + proxy_select_html: proxy_select.html head_insert_html: ui/head_insert.html frame_insert_html: ui/frame_insert.html diff --git a/pywb/manager/manager.py b/pywb/manager/manager.py index d51bf953..335b1bb3 100644 --- a/pywb/manager/manager.py +++ b/pywb/manager/manager.py @@ -6,12 +6,22 @@ import logging from pywb.utils.loaders import load_yaml_config from pywb.utils.timeutils import timestamp20_now from pywb.warc.cdxindexer import main as cdxindexer_main +from pywb.webapp.pywb_init import DEFAULT_CONFIG + +from distutils.util import strtobool +from pkg_resources import resource_string from argparse import ArgumentParser, RawTextHelpFormatter import heapq import yaml +#============================================================================= +# to allow testing by mocking get_input +def get_input(msg): #pragma: no cover + return raw_input(msg) + + #============================================================================= class CollectionsManager(object): """ This utility is designed to @@ -147,11 +157,113 @@ directory structure expected by pywb if len(v) != 2: raise ValueError(msg) + print('Set {0}={1}'.format(v[0], v[1])) metadata[v[0]] = v[1] with open(metadata_yaml, 'w+b') as fh: fh.write(yaml.dump(metadata, default_flow_style=False)) + def _load_templates_map(self): + defaults = load_yaml_config(DEFAULT_CONFIG) + + # Coll Templates + templates = defaults['paths']['template_files'] + + for name, _ in templates.iteritems(): + templates[name] = defaults[name] + + + # Shared Templates + shared_templates = defaults['paths']['shared_template_files'] + + for name, _ in shared_templates.iteritems(): + shared_templates[name] = defaults[name] + + return templates, shared_templates + + def list_templates(self): + templates, shared_templates = self._load_templates_map() + + print('Shared Templates') + for n, v in shared_templates.iteritems(): + print('- {0}: (pywb/{1})'.format(n, v)) + + print('') + + print('Collection Templates') + for n, v in templates.iteritems(): + print('- {0}: (pywb/{1})'.format(n, v)) + + def _confirm_overwrite(self, full_path, msg): + if not os.path.isfile(full_path): + return True + + res = get_input(msg) + try: + res = strtobool(res) + except ValueError: + res = False + + if not res: + raise IOError('Skipping, {0} already exists'.format(full_path)) + + def _get_template_path(self, template_name, verb): + templates, shared_templates = self._load_templates_map() + + try: + filename = templates[template_name] + if not self.coll_name: + msg = ('To {1} a "{0}" template, you must specify ' + + 'a collection name: template --{1} {0}') + raise IOError(msg.format(template_name, verb)) + + full_path = os.path.join(self.templates_dir, + os.path.basename(filename)) + + except KeyError: + try: + filename = shared_templates[template_name] + full_path = os.path.join(os.getcwd(), + os.path.basename(filename)) + except KeyError: + msg = 'template name must be one of {0} or {1}' + msg.format(templates.keys(), shared_templates.keys()) + raise KeyError(msg) + + return full_path, filename + + def add_template(self, template_name, force=False): + full_path, filename = self._get_template_path(template_name, 'add') + + msg = ('Template file "{0}" ({1}) already exists. ' + + 'Overwrite with default template? (y/n) ') + msg = msg.format(full_path, template_name) + + if not force: + self._confirm_overwrite(full_path, msg) + + data = resource_string('pywb', filename) + with open(full_path, 'w+b') as fh: + fh.write(data) + + full_path = os.path.abspath(full_path) + print('Copied default template "{0}" to "{1}"'.format(filename, full_path)) + + def remove_template(self, template_name, force=False): + full_path, filename = self._get_template_path(template_name, 'remove') + + if not os.path.isfile(full_path): + msg = 'Template "{0}" does not exist.' + raise IOError(msg.format(full_path)) + + msg = 'Delete template file "{0}" ({1})? (y/n) ' + msg = msg.format(full_path, template_name) + + if not force: + self._confirm_overwrite(full_path, msg) + + os.remove(full_path) + print('Removed template file "{0}"'.format(full_path)) #============================================================================= def main(args=None): @@ -232,9 +344,37 @@ Create manage file based web archive collections metadata.add_argument('--set', nargs='+') metadata.set_defaults(func=do_metadata) + # Add default template + def do_add_template(r): + m = CollectionsManager(r.coll_name) + if r.add: + m.add_template(r.add, r.force) + elif r.remove: + m.remove_template(r.remove, r.force) + elif r.list: + m.list_templates() + + template_help = 'Add default html template for customization' + template = subparsers.add_parser('template') + template.add_argument('coll_name', nargs='?', default='') + template.add_argument('-f', '--force', action='store_true') + template.add_argument('--add') + template.add_argument('--remove') + template.add_argument('--list', action='store_true') + template.set_defaults(func=do_add_template) + r = parser.parse_args(args=args) r.func(r) +# special wrapper for cli to avoid printing stack trace +def main_wrap_exc(): #pragma: no cover + try: + main() + except Exception as e: + print('Error: ' + str(e)) + sys.exit(2) + + if __name__ == "__main__": - main() + main_wrap_exc() diff --git a/setup.py b/setup.py index 52714325..3783c4b7 100755 --- a/setup.py +++ b/setup.py @@ -90,7 +90,7 @@ setup( cdx-indexer = pywb.warc.cdxindexer:main live-rewrite-server = pywb.apps.live_rewrite_server:main proxy-cert-auth = pywb.framework.certauth:main - wb-manager = pywb.manager.manager:main + wb-manager = pywb.manager.manager:main_wrap_exc """, classifiers=[ 'Development Status :: 4 - Beta', diff --git a/tests/test_auto_colls.py b/tests/test_auto_colls.py index 5d0700fe..a293508e 100644 --- a/tests/test_auto_colls.py +++ b/tests/test_auto_colls.py @@ -14,6 +14,7 @@ from pywb import get_test_dir from pywb.framework.wsgi_wrappers import init_app from pytest import raises +from mock import patch #============================================================================= @@ -39,6 +40,14 @@ def teardown_module(): os.chdir(orig_cwd) +#============================================================================= +mock_input_value = '' + +def mock_raw_input(*args): + global mock_input_value + return mock_input_value + + #============================================================================= class TestManagedColls(object): def setup(self): @@ -266,6 +275,59 @@ class TestManagedColls(object): resp = self.testapp.get('/test/20140103030321/http://example.com?example=1') assert resp.status_int == 200 + def test_add_default_templates(self): + """ Test add default templates: shared, collection, + and overwrite collection template + """ + # list + main(['template', 'foo', '--list']) + + # Add shared template + main(['template', '--add', 'home_html']) + assert os.path.isfile(os.path.join(self.root_dir, 'index.html')) + + # Add collection template + main(['template', 'foo', '--add', 'query_html']) + assert os.path.isfile(os.path.join(self.root_dir, 'collections', 'foo', 'templates', 'query.html')) + + # overwrite -- force + main(['template', 'foo', '--add', 'query_html', '-f']) + + @patch('pywb.manager.manager.get_input', lambda x: 'y') + def test_add_template_input_yes(self): + """ Test answer 'yes' to overwrite + """ + mock_raw_input_value = 'y' + main(['template', 'foo', '--add', 'query_html']) + + + @patch('pywb.manager.manager.get_input', lambda x: 'n') + def test_add_template_input_no(self): + """ Test answer 'no' to overwrite + """ + with raises(IOError): + main(['template', 'foo', '--add', 'query_html']) + + @patch('pywb.manager.manager.get_input', lambda x: 'other') + def test_add_template_input_other(self): + """ Test answer 'other' to overwrite + """ + with raises(IOError): + main(['template', 'foo', '--add', 'query_html']) + + @patch('pywb.manager.manager.get_input', lambda x: 'no') + def test_remove_not_confirm(self): + """ Test answer 'no' to remove + """ + # don't remove -- not confirmed + with raises(IOError): + main(['template', 'foo', '--remove', 'query_html']) + + @patch('pywb.manager.manager.get_input', lambda x: 'yes') + def test_remove_confirm(self): + # remove -- confirm + main(['template', 'foo', '--remove', 'query_html']) + def test_no_templates(self): """ Test removing templates dir, using default template again """ @@ -284,8 +346,11 @@ class TestManagedColls(object): orig_stdout = sys.stdout buff = BytesIO() sys.stdout = buff - main(['list']) - sys.stdout = orig_stdout + + try: + main(['list']) + finally: + sys.stdout = orig_stdout output = sorted(buff.getvalue().splitlines()) assert len(output) == 4 @@ -294,6 +359,23 @@ class TestManagedColls(object): assert '- nested' in output assert '- test' in output + def test_err_template_remove(self): + """ Test various error conditions for templates: + invalid template name, no collection for collection template + no template file found + """ + # no such template + with raises(KeyError): + main(['template', 'foo', '--remove', 'blah_html']) + + # collection needed + with raises(IOError): + main(['template', '--remove', 'query_html']) + + # already removed + with raises(IOError): + main(['template', 'foo', '--remove', 'query_html']) + def test_err_no_such_coll(self): """ Test error adding warc to non-existant collection """