#!/usr/bin/env python3

#   Copyright (C) 2018      Simone Margaritelli
#                 2018      MiWCryptAnalytics
#                 2023      munix9
#                 2023      Wojtek Widomski
#                 2019-2023 Gustavo Iñiguez Goia
#
#   This file is part of OpenSnitch.
#
#   OpenSnitch is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   OpenSnitch is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with OpenSnitch.  If not, see <http://www.gnu.org/licenses/>.

from PyQt5 import QtWidgets, QtCore
from PyQt5.QtNetwork import QLocalServer, QLocalSocket

import sys
import os
import signal
import argparse
import logging

from concurrent import futures

import grpc

dist_path = '/usr/lib/python3/dist-packages/'
if dist_path not in sys.path:
    sys.path.append(dist_path)

from opensnitch.service import UIService
from opensnitch.config import Config
from opensnitch.utils import Themes, Utils, Versions, Message
from opensnitch.utils.xdg import xdg_opensnitch_dir, xdg_current_session

from opensnitch import auth
import opensnitch.proto as proto
ui_pb2, ui_pb2_grpc = proto.import_()

app_id = os.path.join(xdg_opensnitch_dir, "io.github.evilsocket.opensnitch")

def on_exit():
    server.stop(0)
    app.quit()
    try:
        os.remove(app_id)
    except:
        pass
    sys.exit(0)

def restrict_socket_perms(socket):
    """Restrict socket reading to the current user"""
    try:
        if socket.startswith("unix://") and os.path.exists(socket[7:]):
            os.chmod(socket[7:], 0o640)
    except Exception as e:
        print("Unable to change unix socket permissions:", socket, e)

def configure_screen_scale_factor(cfg):
    """configure qt screen scale:
        https://doc.qt.io/qt-5/highdpi.html#high-dpi-support-in-qt
    """
    auto_screen_factor = cfg.getBool(Config.QT_AUTO_SCREEN_SCALE_FACTOR, default_value=True)
    screen_factor = cfg.getSettings(Config.QT_SCREEN_SCALE_FACTOR)
    if screen_factor is None or screen_factor == "":
        screen_factor = "1"

    print("QT_AUTO_SCREEN_SCALE_FACTOR:", auto_screen_factor)
    os.environ["QT_AUTO_SCREEN_SCALE_FACTOR"] = str(int(auto_screen_factor))
    if auto_screen_factor is False:
        print("QT_SCREEN_SCALE_FACTORS:", screen_factor)
        os.environ["QT_SCREEN_SCALE_FACTORS"] = screen_factor

def configure_qt_platform_plugin(cfg):
    qt_plugin = cfg.getSettings(Config.QT_PLATFORM_PLUGIN)
    if qt_plugin is None or qt_plugin == "":
        return

    print("QT_QPA_PLATFORM:", qt_plugin)
    os.environ["QT_QPA_PLATFORM"] = qt_plugin

def check_environ():
    if xdg_current_session == "":
        print("""

  Warning: XDG_SESSION_TYPE is not set.
    If there're no icons on the GUI, please, read the following comment:
    https://github.com/evilsocket/opensnitch/discussions/999#discussioncomment-6579273

""")

def supported_qt_version(major, medium, minor):
    q = QtCore.QT_VERSION_STR.split(".")
    return int(q[0]) >= major and int(q[1]) >= medium and int(q[2]) >= minor

if __name__ == '__main__':
    gui_version, grpcversion, protoversion = Versions.get()
    print("\t ~ OpenSnitch GUI -", gui_version, "~")
    print("\tprotobuf:", protoversion, "-", "grpc:", grpcversion)
    print("-" * 50, "\n")

    parser = argparse.ArgumentParser(description='OpenSnitch UI service.', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("--socket", dest="socket", help='''
Path of the unix socket for the gRPC service (https://github.com/grpc/grpc/blob/master/doc/naming.md).
Default: unix:///tmp/osui.sock

Examples:
    - Listening on Unix socket: opensnitch-ui --socket unix:///tmp/osui.sock
        * Use unix:///run/1000/YOUR_USER/opensnitch/osui.sock for better privacy.
    - Listening on port 50051, all interfaces: opensnitch-ui --socket "[::]:50051"
                        ''', metavar="FILE")
    parser.add_argument("--socket-auth", dest="socket_auth", help="Auth type: simple, tls-simple, tls-mutual")
    parser.add_argument("--tls-ca-cert", dest="tls_ca_cert", help="path to the CA cert")
    parser.add_argument("--tls-cert", dest="tls_cert", help="path to the server cert")
    parser.add_argument("--tls-key", dest="tls_key", help="path to the server key")
    parser.add_argument("--max-clients", dest="serverWorkers", default=10, help="Max number of allowed clients (incoming connections).")
    parser.add_argument("--debug", dest="debug", action="store_true", help="Enable debug logs")
    parser.add_argument("--debug-grpc", dest="debug_grpc", action="store_true", help="Enable gRPC debug logs")
    parser.add_argument("--background", dest="background", action="store_true", help="Start UI in background even, when tray is not available")

    args = parser.parse_args()

    if args.debug:
        import faulthandler
        faulthandler.enable()

    logging.getLogger().disabled = not args.debug
    cfg = Config.get()
    configure_screen_scale_factor(cfg)
    configure_qt_platform_plugin(cfg)

    if args.debug and args.debug_grpc:
        os.environ["GRPC_TRACE"] = "all"
        os.environ["GRPC_VERBOSITY"] = "debug"

    if supported_qt_version(5,6,0):
        try:
            # NOTE: maybe we also need Qt::AA_UseHighDpiPixmaps
            QtCore.QApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling, True)
        except Exception:
            pass

    service = None

    try:
        Utils.create_socket_dirs()
        app = QtWidgets.QApplication(sys.argv)

        localsocket = QLocalSocket()
        localsocket.connectToServer(app_id)

        if localsocket.waitForConnected():
            raise Exception("GUI already running, opening its window and exiting.")
        else:
            localserver = QLocalServer()
            localserver.setSocketOptions(QLocalServer.UserAccessOption)
            localserver.removeServer(app_id)
            localserver.listen(app_id)

        if hasattr(QtCore.Qt, 'AA_UseHighDpiPixmaps'):
            app.setAttribute(QtCore.Qt.AA_UseHighDpiPixmaps, True)
        thm = Themes.instance()
        thm.load_theme(app)

        if args.socket == None:
            # default
            args.socket = "unix:///tmp/osui.sock"

            addr = cfg.getSettings(Config.DEFAULT_SERVER_ADDR)
            if addr != None and addr != "":
                if addr.startswith("unix://"):
                    if not os.path.exists(os.path.dirname(addr[7:])):
                        print("WARNING: unix socket path does not exist, using unix:///tmp/osui.sock, ", addr)
                    else:
                        args.socket = addr
                else:
                    args.socket = addr

        maxmsglen = cfg.getMaxMsgLength()

        service = UIService(app, on_exit, start_in_bg=args.background)
        check_environ()
        localserver.newConnection.connect(service.OpenWindow)
        # @doc: https://grpc.github.io/grpc/python/grpc.html#server-object
        server = grpc.server(futures.ThreadPoolExecutor(),
                                options=(
                                    # https://github.com/grpc/grpc/blob/master/doc/keepalive.md
                                    # https://grpc.github.io/grpc/core/group__grpc__arg__keys.html
                                    # send keepalive ping every 5 second, default is 2 hours)
                                    ('grpc.keepalive_time_ms', 5000),
                                    # after 5s of inactivity, wait 20s and close the connection if
                                    # there's no response.
                                    ('grpc.keepalive_timeout_ms', 20000),
                                    ('grpc.keepalive_permit_without_calls', True),
                                    ('grpc.max_send_message_length', maxmsglen),
                                    ('grpc.max_receive_message_length', maxmsglen),
                                ))

        ui_pb2_grpc.add_UIServicer_to_server(service, server)

        auth_type = auth.Simple
        if args.socket_auth != None:
            auth_type = args.socket_auth
        elif cfg.getSettings(Config.AUTH_TYPE) != None:
            auth_type = cfg.getSettings(Config.AUTH_TYPE)

        # grpc python doesn't seem to accept unix:@address to listen on an
        # abstract unix socket, so use unix-abstract: and transform it to what
        # the Go client understands.
        if args.socket.startswith("unix:@"):
            parts = args.socket.split("@")
            args.socket = "unix-abstract:{0}".format(parts[1])

        print("Using server address:", args.socket, "auth type:", auth_type)

        if auth_type == auth.Simple or auth_type == "":
            server.add_insecure_port(args.socket)
        else:
            auth_ca_cert = args.tls_ca_cert
            auth_cert = args.tls_cert
            auth_certkey = args.tls_key
            if auth_cert == None:
                auth_cert = cfg.getSettings(Config.AUTH_CERT)
            if auth_certkey == None:
                auth_certkey = cfg.getSettings(Config.AUTH_CERTKEY)
            if auth_ca_cert == None:
                auth_ca_cert = cfg.getSettings(Config.AUTH_CA_CERT)

            tls_creds = auth.get_tls_credentials(auth_ca_cert, auth_cert, auth_certkey)
            if tls_creds == None:
                raise Exception("Invalid TLS credentials. Review the server key and cert files.")
            server.add_secure_port(args.socket, tls_creds)

        # https://stackoverflow.com/questions/5160577/ctrl-c-doesnt-work-with-pyqt
        signal.signal(signal.SIGINT, signal.SIG_DFL)

        # print "OpenSnitch UI service running on %s ..." % socket
        server.start()

        restrict_socket_perms(args.socket)

        app.exec_()

    except KeyboardInterrupt:
        on_exit()
    except Exception as e:
        print(e)
    finally:
        if service:
            # finish gracefully, closing notifications channel.
            service.close()
