hpss-agent

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright © 2017-2018 Martin Ueding <martin-ueding.de>

import argparse
import ftplib
import getpass
import json
import logging
import os
import sys


CONFIGPATH = os.path.expanduser('~/.agent-config.js')
logger = logging.getLogger(__name__)


def get_remote_size(ftp, filename):
    '''
    Retrieves the file size in bytes of the remote file.

    :param ftplib.FTP ftp: FTP connection object
    :param str filename: Filename of a file in the current working directory on
        the remote.
    '''
    for path, properties in ftp.mlsd():
        logger.debug('Found %s on archive.', path)
        if os.path.basename(path) == filename:
            return int(properties['size'])

    raise RuntimeError('Remote file not found.')


def assert_correct_transfer(ftp, filename):
    '''
    Asserts that the transfer has been done correctly by requiring exact
    equality of the file sizes locally and removely.

    :param ftplib.FTP ftp: FTP connection object
    :param str filename: Name of the file, file has to be found in the current
         working directory locally and remotely.
    '''
    local_size = os.path.getsize(filename)
    remote_size = get_remote_size(ftp, filename)

    logger.debug('Local size of file is   %d.', local_size)
    logger.debug('Archive size of file is %d.', remote_size)

    assert local_size == remote_size, 'The file in the archive must have the exact same size as on disk.'


def start_session(ftp):
    '''
    :param ftplib.FTP ftp: FTP connection object where the ``connect`` method
        has not been called yet.
    :returns ftplib.FTP: FTP connection object that is now connected and logged
        in.
    '''
    with open(CONFIGPATH) as f:
        config = json.load(f)

    ftp.connect('hpsscore.hww.de', 4021)
    ftp.login(config['username'], config['password'])


def set_class_of_service(ftp, size):
    '''
    The HPSS has a non-standard “class of service” concept. One needs to set it
    before a transfer, otherwise the connection might be dropped prematurely.

    :param ftplib.FTP ftp: FTP connection object
    :param int size: File size in bytes
    '''
    if size < 500 * 10**6:
        cos = 102
    elif size < 8 * 10**9:
        cos = 122
    else:
        cos = 132

    logger.debug('Setting class of service to %d.', cos)

    ftp.sendcmd('SITE setcos {}'.format(cos))


def main_setup():
    username = input('Your HPSS username: ')
    password = getpass.getpass('Your HPSS password: ')

    config = dict(
        username=username,
        password=password,
    )

    with open(CONFIGPATH, 'w') as f:
        json.dump(config, f, indent=2)

    os.chmod(CONFIGPATH, 0o600)


def main_get():
    with ftplib.FTP() as ftp:
        start_session(ftp)

        source_base = os.path.basename(options.source)

        ftp.cwd(os.path.dirname(options.source))
        set_class_of_service(ftp, get_remove_size(ftp, source_base))

        with open(options.source, 'wb') as f:
            logger.info('Starting transfer from archive of %s', source_base)
            ftp.retrbinary('RETR {}'.format(source_base), f.write)

        assert_correct_transfer(ftp, source_base)


def main_put():
    with ftplib.FTP() as ftp:
        start_session(ftp)

        source_base = os.path.basename(options.source)

        if options.destination is not None:
            logger.info('Entering directory %s', options.destination)
            ftp.cwd(options.destination)

        set_class_of_service(ftp, os.path.getsize(options.source))

        with open(options.source, 'rb') as f:
            logger.info('Starting transfer to archive of %s', source_base)
            ftp.storbinary('STOR {}'.format(source_base), f)

        assert_correct_transfer(ftp, source_base)


def main():
    print('Please call with --help.')
    sys.exit(1)


def _parse_args():
    '''
    Parses the command line arguments.

    :return: Namespace with arguments.
    :rtype: Namespace
    '''
    parser = argparse.ArgumentParser(
        description='''
        Wrapper for pftp_client.

        You can use this script to perform the actions listed below. Call this
        script with an action name and `-h` again to get help for that
        particular action.
        ''')
    parser.set_defaults(func=main)
    parser.add_argument('-v', '--verbose', action='count')

    subparsers = parser.add_subparsers(title='Commands')

    parser_get = subparsers.add_parser('get', help='Retrieve file from archive')
    parser_get.add_argument('source')
    parser_get.set_defaults(func=main_get)

    parser_get = subparsers.add_parser('put', help='Put file into archive')
    parser_get.add_argument('source')
    parser_get.add_argument('destination', nargs='?', help='Destination directory in the archive')
    parser_get.set_defaults(func=main_put)

    parser_setup = subparsers.add_parser('setup', help='Setup password file')
    parser_setup.set_defaults(func=main_setup)

    global options
    options = parser.parse_args()

    if options.verbose == 1:
        logging.basicConfig(level=logging.INFO)
    elif options.verbose == 2:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.WARN)

    options.func()


if __name__ == '__main__':
    _parse_args()