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 uuid import uuid4

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_current_redis, get_next_redis
from inventory_provider.tasks.worker import InventoryTask, \
    log_task_entry_and_exit

environment.setup_logging()

logger = logging.getLogger(__name__)


@app.task(base=InventoryTask, bind=True, name='update_equipment_locations_ims')
@log_task_entry_and_exit
def update_equipment_locations_ims(self, use_current=False):

    if use_current:
        r = get_current_redis(InventoryTask.config)
        # scan with bigger batches, to mitigate network latency effects
    else:
        r = get_next_redis(InventoryTask.config)
    rp = r.pipeline()
    for k in r.scan_iter('ims:location:*', count=1000):
        rp.delete(k)
    rp.execute()

    c = InventoryTask.config["ims"]
    ds = IMS(c['api'], c['username'], c['password'])

    rp = r.pipeline()
    hostnames_found = set()
    for h, d in ims_data.get_node_locations(ds):
        # put into a list to match non-IMS version
        rp.set(f'ims:location:{h}', json.dumps([d]))
        if h in hostnames_found:
            print(f'Multiple entries for {h}')
        hostnames_found.add(h)
    rp.execute()


@app.task(base=InventoryTask, bind=True, name='update_lg_routers_ims')
@log_task_entry_and_exit
def update_lg_routers_ims(self, use_current=False):

    if use_current:
        r = get_current_redis(InventoryTask.config)
        for k in r.scan_iter('classifier-cache:ims-lg:*'):
            r.delete(k)
    else:
        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))


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


@app.task(base=InventoryTask, bind=True, name='export_data_for_otrs')
@log_task_entry_and_exit
def export_data_for_otrs(self, files_to_export=None, export_duplicates=False):
    debug_uuid = uuid4()
    logger.debug(f'debug uuid: {debug_uuid}')
    if files_to_export:
        files_to_export = OTRSFiles(files_to_export)
    else:
        files_to_export = set(OTRSFiles)

    ims_config = InventoryTask.config["ims"]
    otrs_config = InventoryTask.config["otrs-export"]

    command_template = 'rsync -aPq --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, return_duplicates=export_duplicates))

        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)
    return debug_uuid