upload.py

#!/usr/bin/env python3

from   datetime  import datetime
from   enum      import Enum
from   mechanize import Browser, _response
from   mimetypes import MimeTypes
from   pathlib   import Path
from   typing    import Dict, Iterable, List, NamedTuple, Optional
import argparse
import json
import mechanize
import re
import subprocess
import urllib
import yaml

################################################################################

class Args(NamedTuple):
    module_dir: Path
    config: Path
    secrets: Path
    server: str
    rm: bool
    replace: bool
    skip_data: bool
    add_date: bool

class Neediness(Enum):
    SKIP = 1
    WANT = 2
    NEED = 3

    def __ge__(self, other: 'Neediness') -> bool:
        return self.value >= other.value

class File(NamedTuple):
    form_name: str
    path: Path
    neediness: Neediness = Neediness.NEED


class Config(NamedTuple):
    users: Optional[List[str]]
    project_module: Dict[str, str]


class Secrets(NamedTuple):
    username: str
    password: str

################################################################################

def read_args() -> Args:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter,
        description='Upload autogenerated module to faims server.'
        "\n"
        "\nThis script uploads a module to the FAIMS server. When run without arguments, it looks for a module in "
        "a directory called 'module', in the same directory as the script. The module's definition files "
        "must have the same names as they would if they were generated by the autogen. These names are as follows:"
        "\n"
        "\n| File purpose | File name            |"
        "\n| ------------ | -------------------- |"
        "\n| Arch16n      | english.0.properties |"
        "\n| CSS          | ui_styling.css       |"
        "\n| Data Schema  | data_schema.xml      |"
        "\n| UI Logic     | ui_logic.bsh         |"
        "\n| UI Schema    | ui_schema.xml        |"
        "\n| Validation   | validation.xml       |"
        "\n"
        "\nAn optional data.tar.gz file can be included in the same directory "
        "too.")

    parser.add_argument(
        'module_dir',
        nargs='?',
        type=Path,
        default=Path(__file__).parent / Path('module'),
        help="path to module document directory")
    parser.add_argument(
        '--config',
        type=Path,
        default=Path(__file__).parent / Path('upload-config.yaml'),
        help="location of the config file in yaml or json format")
    parser.add_argument(
        '--secrets',
        type=Path,
        default=Path(__file__).parent / Path('upload-secrets.yaml'),
        help="location of the secrets file in yaml or json format")
    parser.add_argument(
        '--server',
        type=str,
        default='dev26',
        help='the subdomain of the fedarch.org server where the module will be '
        'uploaded')
    parser.add_argument(
        '--rm',
        action='store_true',
        help="remove module(s) with same base name")
    parser.add_argument(
        '--replace',
        action='store_true',
        help="remove module(s) with same base name and upload new one")
    parser.add_argument(
        '--skip-data',
        action='store_true',
        help="if this argument is specified, the module's data.tar.gz file "
        "will be uploaded")
    parser.add_argument(
        '--add-date',
        action='store_true',
        help="if this argument is specified, the module name will have today's "
        "date appended")

    args = parser.parse_args()

    return Args(
        module_dir=args.module_dir,
        config=args.config,
        secrets=args.secrets,
        server=args.server,
        rm=args.rm,
        replace=args.replace,
        skip_data=args.skip_data,
        add_date=args.add_date
    )

################################################################################

def get_files(module_dir: Path, skip_data: bool) -> List[File]:
    return [
        File(
            'arch16n',
            module_dir / Path('english.0.properties'),
        ),
        File(
            'css_style',
            module_dir / Path('ui_styling.css'),
        ),
        File(
            'data_schema',
            module_dir / Path('data_schema.xml'),
        ),
        File(
            'ui_logic',
            module_dir / Path('ui_logic.bsh'),
        ),
        File(
            'ui_schema',
            module_dir / Path('ui_schema.xml'),
        ),
        File(
            'validation_schema',
            module_dir / Path('validation.xml'),
        ),
        File(
            'data',
            module_dir / Path('data.tar.gz'),
            neediness=Neediness.SKIP if skip_data else Neediness.WANT
        ),
    ]


def ensure_needed_files(files: Iterable[File]) -> None:
    # Make sure all specification files are in the path'd directory
    for file in files:
        if file.neediness >= Neediness.NEED and not file.path.is_file():
            raise FileNotFoundError(file.path)

################################################################################

def read_serialised(serialised: Path) -> Dict:
    with serialised.open('r') as config_file:
        if serialised.suffix.lower() in ('.yaml', '.yml'):
            return yaml.safe_load(config_file)
        elif serialised.suffix.lower() in ('.json',):
            return json.load(config_file)
        else:
            raise Exception('Unrecognised serialisation format')


def read_config(config_path: Path) -> Config:
    config = read_serialised(config_path)
    return Config(
        users=config.get('users'),
        project_module={
            str(k): str(v) for k, v in config['project_module'].items()
        },
    )


def read_secrets(secrets_path: Path) -> Secrets:
    secrets = read_serialised(secrets_path)
    return Secrets(
        username=secrets['username'],
        password=secrets['password'],
    )

################################################################################

def get_module_name(args: Args, config: Config) -> str:
    base_module_name = config.project_module.get('name')

    if not base_module_name:
        base_module_name = subprocess.run(
            [
                '/usr/bin/env',
                'bash',
                '-c',
                'basename `git rev-parse --show-toplevel`'
            ],
            capture_output=True
        ).stdout.decode('utf-8').strip()

    if not base_module_name:
        base_module_name = args.module_dir.parent.absolute().name

    return base_module_name + (
        datetime.now().strftime(' %Y-%m-%d') if args.add_date else '')

################################################################################

def find_csrf_token(res: mechanize._response.response_seek_wrapper) -> str:
    token = res.get_data()
    token = token.decode('utf-8')
    token = re.search('"([^"]+)"\s+name="csrf-token"', token)
    return token.group(1)

################################################################################

def add_file(br: Browser, filename: Path, input_name: str) -> None:
    br.form.add_file(filename.open('rb'), filename=filename.name, name=input_name)


def add_files(br: Browser, files: Iterable[File]) -> None:
    for file in files:
        add_file(br, file.path, f'project_module[{file.form_name}]')


def login(br: Browser, secrets: Secrets) -> None:
    # Get login form and fill it out
    print('Downloading login form...')
    br.select_form(nr=0)
    br['user[email]'], br['user[password]'] = secrets

    # Submit yer good ol' form
    print('Submitting login form...')
    print()
    res = br.submit()

    if res.geturl().endswith('/sign_in'):
        raise Exception('Could not login (likely due to invalid credentials)')


def go_home(br: Browser) -> None:
    br.follow_link(nr=0)


def delete_module(br: Browser, args: Args, config: Config) -> None:
    module_name = get_module_name(args, config)

    print(f'Deleting module with name {module_name}...')
    # Try clicking on link to module config page
    try:
        res = br.follow_link(text=module_name)
    except:
        print('Cannot delete module; does not exist.')
        print()
        return

    # Get URL to delete module
    req      = br.click_link(text='Delete Module')
    url      = req.get_full_url()

    # Search for the stupid CSRF token
    token    = find_csrf_token(res)
    params   = {
            u'_method'            : 'delete',
            u'authenticity_token' : token
    }
    data = urllib.parse.urlencode(params)

    # POST the deletion request
    try:
        br.open(url, data)
        print('Module deleted')
    except:
        print('Module could not be deleted')

    print()


def set_users(br: Browser, module_name: str, users: Optional[List[str]]) -> None:
    if users is None:
        return

    print('Setting users...')
    # assumes browser is at home page
    link = br.find_link(text=module_name)
    base_link = link.absolute_url.replace(
            '/project_modules/', '/project_module/')

    res = br.open(base_link + '/edit_project_module_user')
    data = res.get_data()
    html = data.decode('utf-8')

    # Get dict(username -> user ID)
    available_users = re.findall(
        pattern=r"<option value='(\d+)'>([^<]+)</option>",
        string=html,
    )
    username_to_user_id = {
        username: int(user_id)
        for user_id, username in available_users}

    # Add users
    users_to_add = {
        username: username_to_user_id[username]
        for username in users
        if username in username_to_user_id}

    for username, user_id in users_to_add.items():
        print(f'Adding user {username}...')
        params = {
            u'authenticity_token' : find_csrf_token(res),
            u'user_id'            : user_id,
        }
        data = urllib.parse.urlencode(params)

        res = br.open(base_link + '/update_project_module_user', data)

    # Get selected users
    data = res.get_data()
    html = data.decode('utf-8')

    selected_fnames = re.findall(
        pattern=r"name='fname\[\]' readonly='readonly' value='([^']+)",
        string=html)
    selected_lnames = re.findall(
        pattern=r"name='lname\[\]' readonly='readonly' value='([^']+)",
        string=html)
    all_deletion_links = br.links(url_regex=r'remove_project_module_user')
    deletion_links_to_click = [
        (f'{fname} {lname}', link.absolute_url)
        for fname, lname, link in zip(selected_fnames, selected_lnames, all_deletion_links)
        if f'{fname} {lname}' not in users
    ]

    for username, link in deletion_links_to_click:
        print(f'Removing user {username}...')
        params = {
            u'authenticity_token' : find_csrf_token(res),
        }
        data = urllib.parse.urlencode(params)
        res = br.open(link, data)

    print()


def upload_module(
        br: Browser,
        args: Args,
        config: Config,
        files: Iterable[File]) -> None:
    module_name = get_module_name(args, config)

    print(f'Uploading module with name {module_name}...')

    # Determine if module's already been uploaded. (Yes, I know I've misused the
    # try-except statement.)
    print('Downloading update/create form...')
    try:
        # The module's already been uploaded
        br.follow_link(text=module_name)
        br.follow_link(text='Edit Module')

        print('Submitting update form...')
        print()
        br.select_form(nr=0)

        form_names = set([
            'arch16n',
            'css_style',
            'ui_logic',
            'ui_schema',
            'validation_schema',
        ])
        filtered_files = (f for f in files if f.form_name in form_names)
        add_files(br, files=filtered_files)

        br.submit()
    except:
        # We need to create a new module
        br.follow_link(text='Create Module')

        print('Submitting create form...')
        print()
        br.select_form(nr=0)
        for k, v in config.project_module.items():
            br[f'project_module[{k}]'] = v
        br[f'project_module[name]'] = module_name

        form_names = set([
            'arch16n',
            'css_style',
            'data_schema',
            'ui_logic',
            'ui_schema',
            'validation_schema',
        ])
        filtered_files = (f for f in files if f.form_name in form_names)
        add_files(br, files=filtered_files)

        failed_msg = 'Failed to create module.'
        res = br.submit()
        data = res.get_data()
        data = data.decode('utf-8')
        if failed_msg in data:
            print(f'Server returned error: {failed_msg}')
            return
        reg = '<span class="help-inline">([^<]+)</span>'
        matches = re.findall(reg, data)
        if matches:
            print('Completed with errors:')
            for m in matches:
                print('  - ', m.group(1))
            print()
            return

        set_users(br, module_name, config.users)

    data_file = next(iter(f for f in files if f.form_name == 'data'), None)
    if (
            data_file and
            data_file.neediness >= Neediness.WANT and
            data_file.path.is_file()):
        print(f'Uploading {data_file.path}...')
        go_home(br)
        br.follow_link(text=module_name)
        br.follow_link(text='Upload Files')
        br.select_form(nr=0)
        add_file(br, data_file.path, 'project_module[file]')
        br.submit()

    print()


################################################################################

def main() -> int:
    args = read_args()
    config = read_config(args.config)
    secrets = read_secrets(args.secrets)

    br = Browser()
    br.open(f'http://{args.server}.fedarch.org')

    files = get_files(args.module_dir, args.skip_data)
    ensure_needed_files(files)

    # Navigate website and upload module
    login(br, secrets)
    if args.replace:
        delete_module(br, args, config)
        go_home      (br)
        upload_module(br, args, config, files)
    elif args.rm:
        delete_module(br, args, config)
    else:
        upload_module(br, args, config, files)

    return 0


if __name__ == '__main__':
    main()