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 pathlib import Path

import click
from decouple import config
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('file_snapshot')
    def run(self, file_snapshot):
        """Run command."""
        self.file_snapshot = file_snapshot
        self.snapshot_json = None
        self.json_wb = None

        with self.app.app_context():
            self.get_user()
            self.open_snapshot()
            self.build_snapshot()

    def get_user(self):
        """Get datamodel of user."""
        self.email = config('EMAIL_DEMO')
        self.password = config('PASSWORD_DEMO')
        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 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.snapshot_json:
            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")