devicehub-teal/ereuse_devicehub/commands/snapshots.py

"""This command is used for up one snapshot."""

import json

# from uuid import uuid4
from io import BytesIO
from os import listdir
from os import remove as remove_file
from os.path import isfile, join
from pathlib import Path

import click
from flask.testing import FlaskClient
from flask_wtf.csrf import generate_csrf

from ereuse_devicehub.resources.user.models import User


class UploadSnapshots:
    """Command.

    This command allow upload all snapshots than exist
    in the directory snapshots_upload.
    If this snapshot exist replace it.

    """

    def __init__(self, app) -> None:
        """Init function."""
        super().__init__()
        self.app = app
        self.schema = app.config.get('DB_SCHEMA')
        self.app.cli.command('snapshot', short_help='Upload snapshots.')(self.run)

    @click.argument('email')
    @click.argument('password')
    def run(self, email, password=None):
        """Run command."""
        self.email = email
        self.password = password
        self.json_wb = None
        self.onlyfiles = []

        self.get_user()
        self.get_files()
        for f in self.onlyfiles:
            self.file_snapshot = f
            self.open_snapshot()
            self.build_snapshot()
        self.remove_files()

    def get_user(self):
        """Get datamodel of user."""
        self.user = User.query.filter_by(email=self.email).one()
        self.client = FlaskClient(self.app, use_cookies=True)
        self.client.get('/login/')

        data = {
            'email': self.email,
            'password': self.password,
            'remember': False,
            'csrf_token': generate_csrf(),
        }
        self.client.post('/login/', data=data, follow_redirects=True)

    def remove_files(self):
        """Open snapshot file."""
        for f in self.onlyfiles:
            remove_file(Path(__file__).parent.joinpath('snapshot_files').joinpath(f))

    def open_snapshot(self):
        """Open snapshot file."""
        with Path(__file__).parent.joinpath('snapshot_files').joinpath(
            self.file_snapshot,
        ).open() as file_snapshot:
            self.json_wb = json.loads(file_snapshot.read())
            b_snapshot = bytes(json.dumps(self.json_wb), 'utf-8')
            self.file_snap = (BytesIO(b_snapshot), self.file_snapshot)

    def build_snapshot(self):
        """Build the devices of snapshot."""
        uri = '/inventory/upload-snapshot/'

        if not self.json_wb:
            return

        self.client.get(uri)
        data = {
            'snapshot': self.file_snap,
            'csrf_token': generate_csrf(),
        }

        self.client.post(uri, data=data, content_type="multipart/form-data")

    def get_files(self):
        """Read snaoshot_files dir."""
        mypath = Path(__file__).parent.joinpath('snapshot_files')
        for f in listdir(mypath):
            if not isfile(join(mypath, f)):
                continue
            if not f[-5:] == ".json":
                continue
            self.onlyfiles.append(f)