munyal-client/client.py

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

import base64
import json
import os
import pathlib
import sys
import shutil
from copy import deepcopy
from datetime import datetime
from random import randint
from shutil import rmtree
from threading import Thread
from time import sleep
from uuid import uuid4

from PIL import Image, ImageTk
from requests import post

import websocket
from compare_json import compare_json
from config import get_config
from dir_to_json import get_json
from misc import check_network
from misc import flatten_dirs
from misc import path_join
from pystray import Icon, Menu
from pystray import MenuItem as Item
from websocket import WebSocket

try:
    import thread
except ImportError:
    import _thread as thread

_online_icon = Image.open("img/icons/online.png")
_offline_icon = Image.open("img/icons/offline.png")
_standby_icon = Image.open("img/icons/offline.png")

if not int(os.getenv("VERBOSE", 1)):
    def X(*args, **kwargs):
        pass
    print = X

class MunyalClient(Icon):
    def __init__(self):
        super(MunyalClient, self).__init__("Munyal")
        self.__ws_break = False
        self.ws = None
        self.icon = _standby_icon
        self.config = get_config()
        self.stack = []
        self.ignored = []

    def is_online(self):
        # if not self.ws_online:
        #     self.icon = _standby_icon
        #     return False
        ping = check_network("http://google.com", 443)
        if ping:
            self.icon = _online_icon
        else:
            self.icon = _offline_icon
        print(f"La PC está {'on' if bool(ping) else 'off'}line")
        return ping

    def __ws_message(self, ws, message):
        print(message)

    def __ws_error(self, ws, error):
        print("Error: ", error)

    def __ws_open(self):
        print("WebSocket abierto")
        thread.start_new_thread(self.listener, ())

    def start(self):
        websocket.enableTrace(True)
        # ws.on_open = self.listener
        thread_uploader = Thread(target=self.uploader,
                                 name="uploader",
                                 daemon=True)
        thread_downloader = Thread(target=self.downloader,
                                   name="downloader",
                                   daemon=True)
        # thread_listener = Thread(target=self.listener,
        #                          name="listener",
        #                          daemon=True)
        thread_uploader.start()
        # thread_downloader.start()
        # thread_listener.start()

        self.run(self.__run)
        self.menu = Menu(Item("Exit", lambda *args: sys.exit(0)))

    def __run(self, icon):
        icon.visible = True
        while True:
            print("Conectando al websocket")
            ws = websocket.WebSocketApp(
                f"ws://{self.config['login']['user']}.loca.lt",
                # "ws://127.0.0.1:12345",
                on_message=self.__download,
                on_error=print,
                on_close=lambda soc: self.__ws_close(ws))
            ws.on_open = lambda soc: self.listener(soc)
            ws.run_forever()
            sleep(1)

    def __ws_close(self, ws):
        print("WebSocket cerrado")
        self.__ws_break = True

    def listener(self, ws):
        self.ws = ws

        def wrapper(*args):
            folder = self.config["folder"]
            uuid = self.config["uuid"]
            while True:
                print(self.stack)
                if self.stack:
                    try:
                        data = self.stack[0]
                        data["uuid"] = uuid
                        if data["name"] in self.ignored:
                            self.ignored.remove(data["name"])
                            self.stack.pop(0)
                            continue
                        if data["is_file"] and data["action"] == "add":
                            full_path = path_join(folder, data["name"])
                            data["file"] = base64.b85encode(
                                open(full_path, "rb").read()).decode("ascii")
                        self.ws.send(json.dumps(data))
                        self.stack.pop(0)
                    except Exception as e:
                        print(e)
                        print("Error uploading file, trying again")
                if not self.is_online() or self.__ws_break:
                    self.__ws_break = False
                    self.ws = None
                    return
                sleep(1)

        thread.start_new_thread(wrapper, ())

    def uploader(self):
        print("Uploader")
        folder = self.config["folder"]
        last = get_json(folder)
        while True:
            sleep(1)
            actual = get_json(folder)
            _actual = deepcopy(actual)
            delete, add = compare_json(last, actual)

            if delete:
                print("Cosas eliminadas:")
                print(delete)
                delete = flatten_dirs(delete, action="delete")
                for f in delete:
                    f["action"] = "delete"
                self.stack.extend(delete)
            if add:
                print("Cosas agregadas:")
                print(add)
                add = flatten_dirs(add, action="add")
                for f in add:
                    f["action"] = "add"
                self.stack.extend(add)
            last = _actual

    def __upload(self, path):
        pass

    def downloader(self, data):
        name = data["name"]
        self.ignored.append(name)
        full_path = path_join(self.config["folder"], name)
        if data["is_file"]:
            if data["action"] == "add":
                directory = os.path.split(full_path)[0]
                if not os.path.exists(directory):
                    os.makedirs(directory, exist_ok=True)
                directory, _ = os.path.split(full_path)
                temp_name = str(uuid4()) + ".tmp"
                temp_dir = path_join(directory, temp_name)
                with open(temp_dir, "wb") as f:
                    f.write(base64.b85decode(data["file"].encode("ascii")))
                os.rename(temp_dir, full_path)
            else:
                os.remove(full_path)
        else:
            if data["action"] == "add":
                os.mkdir(full_path)
            else:
                shutil.rmtree(full_path)

    def __download(self, message):
        try:
            data = json.loads(message)
        except Exception:
            return
        if data["uuid"] != self.config["uuid"]:
            self.downloader(data)


if __name__ == '__main__':
    try:
        client = MunyalClient()
        client.start()
        sys.exit(0)
    except BaseException as e:
        print(e)
        sys.exit(1)