munyal-client/munyal/client.py

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

import base64
import json
import os
import pathlib
import sys
import shutil
import websocket

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

from pystray import Icon, Menu
from pystray import MenuItem as Item
from websocket import WebSocket

from munyal.compare_json import compare_json
from munyal.config import get_config
from munyal.config import get_last_history
from munyal.config import set_last_history
from munyal.dir_to_json import get_json
from munyal.misc import check_network
from munyal.misc import path_join


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/standby.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
        if self.__ws_break:
            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.menu = Menu(Item("Exit", lambda *args: sys.exit(0)))
        self.run(self.__run)

    def __run(self, icon):
        icon.visible = True
        while True:
            print("Conectando al websocket")
            websocket.enableTrace(False)
            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
        self.icon = _standby_icon

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

        def wrapper(*args):
            folder = self.config["folder"]
            uuid = self.config["uuid"]
            while True:
                if self.stack:
                    data = self.stack[0]
                    try:
                        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:
                        if data["is_file"] and data["action"] == "add":
                            path = path_join(folder, data["name"])
                            if not os.path.exists(path):
                                self.stack.pop(0)
                            print("the file does not exists anymore, skiped")
                        else:
                            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_last_history()
        if last is None:
            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(len(delete))
                for f in delete:
                    f["action"] = "delete"
                self.stack.extend(delete)
            if add:
                print("Cosas agregadas:")
                print(len(add))
                for f in add:
                    f["action"] = "add"
                self.stack.extend(add)
            last = _actual
            if delete or add:
                set_last_history(last)

    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:
                while True:
                    try:
                        os.remove(full_path)
                        break
                    except:
                        print("Maybe the file is open, please close first")
                        sleep(5)
        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)