#!/usr/bin/python

"""
Merges add_servers into current favorites and removes remove_servers.
Run as root to update all users or as normal user to update just that user.
"""

import os
import getpass
import subprocess
import uuid

import Foundation

favorites_path = "/Users/{user}/Library/Application Support/com.apple.sharedfilelist/com.apple.LSSharedFileList.FavoriteServers.sfl2"

# add these servers if they don't exist already
# use a tuple: ("<name>", "<path>") to set a name for the favorite.
# otherwise just use a string and the path will be used as the name
add_servers = (("My Name", "smb://server.example.com/share"), "vnc://server.example.com")
# remove these servers if they exist. Use a wildcard (*) at the end to remove all shares
# does not support ("<name>", "<path>") syntax
remove_servers = ("smb://old.example.com/*", "vnc://old.example.com")

def get_users():
    "Get users with a home directory in /Users"

    # get users from dscl
    dscl_users = subprocess.check_output(["/usr/bin/dscl", ".", "-list", "/Users"]).splitlines()

    # get home directories
    homedir_users = os.listdir("/Users")

    # return users that are in both lists
    users = set(dscl_users).intersection(set(homedir_users))
    return [u.strip() for u in users if u.strip() != ""]

def set_favorites(user, add_servers, remove_servers):
    "Set the Server Favorites for the given user"

    # read existing favorites file
    data = Foundation.NSKeyedUnarchiver.unarchiveObjectWithFile_(favorites_path.format(user=user))

    # reformat add_servers to [(name, path), ...]
    new_add_servers = []
    for s in add_servers:
        if len(s) == 2:
            new_add_servers.append(s)
        else:
            new_add_servers.append((s, s))
    add_servers = new_add_servers

    existing_servers = []

    # read existing items safely
    if data is not None:
        data_items = data.get("items", [])

        # read existing servers
        for item in data_items:
            name = item["Name"]
            url, _, _ = Foundation.NSURL.initByResolvingBookmarkData_options_relativeToURL_bookmarkDataIsStale_error_(Foundation.NSURL.alloc(), item["Bookmark"], Foundation.NSURLBookmarkResolutionWithoutUI, None, None, None)
            existing_servers.append((name, unicode(url)))

    # get unique ordered list of all servers
    all_servers = []

    # add existing_servers to list, updating name if necessary
    for s in existing_servers:
        try:
            idx = [a[1] for a in add_servers].index(s[1])
            all_servers.append((add_servers[idx][0], s[1]))
        except ValueError:
            all_servers.append(s)

    # add add_servers if not present already
    for s in add_servers:
        if s[1] not in [e[1] for e in existing_servers]:
            all_servers.append(s)

    # remove old servers: exact match
    # matches "smb://old.domain" exactly
    all_servers = [s for s in all_servers if s[1] not in remove_servers]

    # remove old servers: shares
    # matches "smb://old.domain/*"
    all_servers = [s for s in all_servers if len(
        [True for r in remove_servers if r.endswith("*") and s[1].startswith(r[:-1])]
    ) < 1]

    # generate necessary structures
    items = []
    for name, path in all_servers:
        item = {}
        # use unicode to translate to NSString
        item["Name"] = unicode(name)
        url = Foundation.NSURL.URLWithString_(unicode(path))
        bookmark, _ = url.bookmarkDataWithOptions_includingResourceValuesForKeys_relativeToURL_error_(0, None, None, None)
        item["Bookmark"] = bookmark
        # generate a new UUID for each server
        item["uuid"] = unicode(uuid.uuid1()).upper()
        item["visibility"] = 0
        item["CustomItemProperties"] = Foundation.NSDictionary.new()

        items.append(Foundation.NSDictionary.dictionaryWithDictionary_(item))

    data = Foundation.NSDictionary.dictionaryWithDictionary_({
        "items": Foundation.NSArray.arrayWithArray_(items),
        "properties": Foundation.NSDictionary.dictionaryWithDictionary_({"com.apple.LSSharedFileList.ForceTemplateIcons": False})
    })

    # write sfl2 file
    Foundation.NSKeyedArchiver.archiveRootObject_toFile_(data, favorites_path.format(user=user))

# loop through users and set favorites
if __name__ == "__main__":
    # if running as root, run for all users. Otherwise run for current user
    user = getpass.getuser()
    if user == "root":
        users = get_users()
    else:
        users = [user]

    for user in users:
        try:
            set_favorites(user, add_servers, remove_servers)
            # fix owner if ran as root
            if user == "root":
                os.system(("chown {user} " + favorites_path).format(user=user))
            print "Server Favorites set for " + user
        except Exception as e:
            # if there's an error, log it an continue on
            print "Failed setting Server Favorites for {0}: {1}".format(user, str(e))

    # kill sharedfilelistd process to reload file. Finder should be closed when this happens
    os.system("killall sharedfilelistd")