#!/usr/bin/env python3
# rdp -- launch FreeRDP to make Remote Desktop connections
import argparse
import getpass
import gi
import logging
import os
import socket
import subprocess
import sys

gi.require_version("Secret", "1")
from gi.repository import Secret

def die(msg):
    if sys.stdin.isatty():
        exit("error: %s" % msg)
    else:
        msg = msg.replace("&", "&amp;") \
                 .replace("<", "&lt;") \
                 .replace(">", "&gt;")
        subprocess.run(["zenity", "--error",
                                  "--title=mstsc",
                                  "--text=<b>Error:</b> %s" % msg,
                                  "--width=300"])
        exit(1)

def argv_hide_password(args):
    nargs = []
    for arg in args:
        if arg.startswith("/p:"):
            arg = "/p:" + "*"*(len(arg)-len("/p:"))
        nargs.append(arg)
    return nargs

def argv_join(args):
    nargs = []
    for arg in args:
        if " " in arg:
            arg = "'%s'" % arg
        nargs.append(arg)
    return " ".join(nargs)

def which(cmd):
    dirs = os.environ["PATH"].split(":")
    for d in dirs:
        exe = os.path.join(d or ".", cmd)
        if os.path.exists(exe):
            return exe
    raise KeyError(f"{cmd!r} not found in PATH")

def resolve_fqdn(host):
    res = socket.getaddrinfo(host, 3389,
                             type=socket.SOCK_STREAM,
                             flags=socket.AI_CANONNAME)
    return res[0][3]

def ask_credentials_zenity(host, title):
    if sys.stdout.isatty():
        username = input("Username: ")
        password = getpass.getpass("Password: ")
        return username, password
    else:
        res = subprocess.run(["zenity",
                              "--forms",
                              "--title", "Enter credentials",
                              "--text", "Enter credentials for %s:" % title,
                              "--add-entry", "Username:",
                              "--add-password", "Password:",
                              "--separator", "\n"],
                             stdout=subprocess.PIPE)
        if res.stdout:
            username, password, _ = res.stdout.decode().split("\n")
            return username, password
    return None, None

def ask_credentials(host):
    title = "Remote Desktop on %s" % host
    return ask_credentials_zenity(host, title)

def get_credentials_libsecret(host):
    targets = [host]
    if len(host.split(".")) >= 3:
        _, domain = host.split(".", 1)
        targets = [host, "*." + domain]

    for target in targets:
        for proto in ["rdp", "smb"]:
            qattrs = {"xdg:schema": "org.gnome.keyring.NetworkPassword",
                      "protocol": proto,
                      "server": target}
            logging.debug("searching keyring for protocol %r server %r", proto, target)
            res = Secret.password_search_sync(None,
                                              qattrs,
                                              Secret.SearchFlags.LOAD_SECRETS,
                                              None)
            if res:
                item = res[0]
                logging.debug("found keyring entry %r", item.get_label())
                for k, v in sorted(item.get_attributes().items()):
                    logging.debug(" - %s = %r", k, v)
                username = item.get_attributes()["user"]
                password = item.get_secret().get_text()
                if domain := item.get_attributes().get("domain"):
                    logging.warn("XXX: Found credential with separate domain, u=%r d=%r",
                                 username, domain)
                    if "." in domain:
                        username = "%s@%s" % (username, domain)
                    else:
                        username = "%s\\%s" % (domain, username)
                return username, password

def save_credentials(host, username, password):
    label = "Remote Desktop on %s" % host
    attrs = {"xdg:schema": "org.gnome.keyring.NetworkPassword",
             "protocol": "rdp",
             "server": host,
             "user": username}
    Secret.password_store_sync(None,
                               attrs,
                               Secret.COLLECTION_DEFAULT,
                               label,
                               password,
                               None)

def get_credentials(host, ask=True):
    if r := get_credentials_libsecret(host):
        return r
    if ask:
        print("Credentials not found, asking interactively")
        return ask_credentials(host)
    return None, None

parser = argparse.ArgumentParser()
parser.add_argument("-4", "--ipv4",
                        action="store_true",
                        help="prefer IPv4 for connection")
parser.add_argument("-l", "--legacy", "--old",
                        action="store_true",
                        help="legacy server (Windows 2003 or older, non-NLA)")
parser.add_argument("-r", "--restricted-admin", "--ra",
                        action="store_true",
                        help="Restricted Admin mode - don't delegate any credentials")
parser.add_argument("-k", "--insecure",
                        action="store_true",
                        help="ignore certificates")
parser.add_argument("-W", "--workarea",
                        action="store_true",
                        help="fit display into the work area instead of full-screen")
parser.add_argument("-w", "--no-fullscreen",
                        action="store_true",
                        help="use a non-full-screen window of default size")
parser.add_argument("-g", "--dimensions", metavar="WxH",
                        help="use specified remote screen size")
parser.add_argument("--no-gfx",
                        action="store_true",
                        help="disable support for RDP8 \"Graphics Pipeline\"")
parser.add_argument("-m", "--modem",
                        action="count", default=0,
                        help="set legacy \"network connection type\"")
parser.add_argument("--dev",
                        action="store_true",
                        help="use local FreeRDP build if available")
parser.add_argument("--flatpak",
                        action="store_true",
                        help="use the Flatpak build")
parser.add_argument("--x11",
                        default=True,
                        action="store_true",
                        help="use the X11 client instead of SDL client")
parser.add_argument("--save",
                        action="store_true",
                        help="save password for future connections")
parser.add_argument("-v", "--verbose",
                        action="store_true",
                        help="show more details")
parser.add_argument("host")
args = parser.parse_args()

logging.basicConfig(level=[logging.INFO, logging.DEBUG][args.verbose],
                    format="rdp: %(message)s")

# We need the FQDN for keyring lookup (and kerberos, and possibly certificate
# validation).
try:
    fqdn = resolve_fqdn(args.host)
    logging.debug("resolved host %r to FQDN %r", args.host, fqdn)
except socket.gaierror as e:
    exit("rdp: could not resolve hostname %r: %s" % (args.host, e))

if args.save:
    username, password = ask_credentials(fqdn)
    if username and password:
        save_credentials(fqdn, username, password)
    exit()
else:
    username, password = get_credentials(fqdn, ask=(not args.legacy))

if args.x11:
    cmd = ["xfreerdp"]
else:
    cmd = ["sdl-freerdp"]

if args.flatpak:
    cmd = ["flatpak", "run", "--command=%s" % cmd[0], "com.freerdp.FreeRDP"]
else:
    if args.dev:
        os.environ["PATH"] = "%s:%s" % (os.path.expanduser("~/.local/pkg/FreeRDP/bin"),
                                        os.environ["PATH"])
    try:
        cmd[0] = which(cmd[0] + "3")
    except KeyError:
        cmd[0] = which(cmd[0])

logging.debug("using FreeRDP executable: %s", cmd[0])

cmd.append("/t:Remote Desktop: %s" % args.host)
cmd.append("/v:%s" % fqdn.lower())

if username and password:
    cmd.append("/u:%s" % username)
    cmd.append("/p:%s" % password)
else:
    # Try to avoid the interactive pre-prompt with --legacy
    cmd.append("/u:")
    cmd.append("/p:")

if args.legacy:
    cmd.append("/sec:rdp")
    cmd.append("/cert:ignore")
    if args.restricted_admin:
        die("Restricted Admin mode is not supported by legacy servers")
    #cmd.append("/tls:seclevel:0")
else:
    cmd.append("/ipv4" if args.ipv4 else "/ipv6")
    cmd.append("/sec:nla")
    if args.restricted_admin:
        cmd.append("/restricted-admin")
    if args.insecure:
        cmd.append("/cert:ignore")

cmd.append("/network:auto")

if args.no_gfx:
    cmd.append("/gfx:progressive:off,RFX:off,AVC420:off,AVC444:off")

if args.modem:
    cmd.append("/network:modem")
    cmd.append("/bpp:8" if args.modem >= 2 else "/bpp:16")
else:
    #cmd.append("/bpp:32")
    #cmd.append("+fonts")
    #cmd.append("+aero")
    #cmd.append("+window-drag")
    #cmd.append("+wallpaper")
    #cmd.append("+multitouch")
    cmd.append("/sound:sys:pulse")

cmd.append("+clipboard")
#cmd.append("+home-drive")
#cmd.append("/drive:media,/run/media/%s" % os.getlogin())
#cmd.append("/smartcard")

if args.legacy:
    cmd.append("/cache:bitmap:on,persist")

if args.workarea:
    cmd.append("/workarea")
    #cmd.append("-decorations")
elif args.no_fullscreen:
    cmd.append("/size:%dx%d" % (1280, 896))
elif args.dimensions:
    cmd.append("/size:%s" % args.dimensions)
else:
    cmd.append("/f")
    cmd.append("+floatbar")

if args.verbose:
    cmd.append("/log-level:DEBUG")
else:
    cmd.append("/log-level:WARN")

logging.debug("running with options: %s", argv_join(argv_hide_password(cmd)))
subprocess.run(cmd)
