ims_worker.py

import csv
import json
import logging
import subprocess
import tempfile
from datetime import datetime
from enum import IntFlag
from pathlib import Path

from inventory_provider.db import ims_data
from inventory_provider.db.ims import IMS
from inventory_provider import environment
from inventory_provider.tasks.app import app
from inventory_provider.tasks.common import get_next_redis
from inventory_provider.tasks.worker import InventoryTask

environment.setup_logging()

logger = logging.getLogger(__name__)


@app.task(base=InventoryTask, bind=True)
def update_lg_routers_ims(self):
    logger.debug('>>> update_lg_routers_ims')

    r = get_next_redis(InventoryTask.config)
    for k in r.scan_iter('ims:lg:*'):
        r.delete(k)
    c = InventoryTask.config["ims"]
    ds = IMS(c['api'], c['username'], c['password'])

    for router in ims_data.lookup_lg_routers(ds):
        r.set(f'ims:lg:{router["equipment name"]}', json.dumps(router))

    logger.debug('<<< update_lg_routers_ims')


class OTRSFiles(IntFlag):
    CUSTOMER_COMPANIES = 1
    CUSTOMER_USERS = 2


@app.task(base=InventoryTask, bind=True)
def export_data_for_otrs(self, files_to_export=None):
    if files_to_export:
        files_to_export = OTRSFiles(files_to_export)
    else:
        files_to_export = set(OTRSFiles)

    logger.debug('>>> export_data_for_otrs_ims')
    ims_config = InventoryTask.config["ims"]
    otrs_config = InventoryTask.config["otrs-export"]

    command_template = 'rsync -aP --rsh="ssh -l {user} -p 22 -i {key_file} -o \'UserKnownHostsFile {known_hosts}\'" {source_dir}/* {destination}'  # noqa

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)
        ds = IMS(
            ims_config['api'],
            ims_config['username'],
            ims_config['password'])

        prefix = datetime.now().strftime('%Y%m%d') + '_'

        if OTRSFiles.CUSTOMER_COMPANIES in files_to_export:
            cus_co_path = temp_path.joinpath(f'{prefix}customer_company.csv')
            with open(cus_co_path, 'w+') as f:
                writer = csv.writer(f, delimiter='^')
                writer.writerows(ims_data.otrs_get_customer_company_rows(ds))

        if OTRSFiles.CUSTOMER_USERS in files_to_export:
            cus_usr_path = temp_path.joinpath(f'{prefix}customer_user.csv')
            with open(cus_usr_path, 'w+') as f:
                writer = csv.writer(f, delimiter='^')
                writer.writerows(ims_data.otrs_get_customer_users_rows(ds))

        command = command_template.format(
            user=otrs_config['username'],
            key_file=otrs_config['private-key'],
            known_hosts=otrs_config['known-hosts'],
            source_dir=temp_dir,
            destination=otrs_config['destination']
        )
        subprocess.run(command, shell=True, check=True)

    logger.debug('<<< export_data_for_otrs_ims')