#!/bin/bash -u
#
# wine-import-extensions - Register native file extensions in Wine
#
#    Copyright (C) 2011 Rodrigo Silva (MestreLion) <linux@rodrigosilva.com>
#
#    This program 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.
#
#    This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Huge thanks for all the gurus and friends in irc://irc.freenet.org/#bash
# and the contributors of http://mywiki.wooledge.org/
#
# FIXME: Properly handle cases in exts and mime-types. Windows might be case-
#  insensitive but winemenubuilder is not. And it is buggy about it.
# TODO: undo option (from changes file)
# TODO: --include-only and --force-all options
# TODO: update list with status for each extension (new, updated, skipped, etc)
# TODO: handle URIs (mailto:, magnet:, irc://)


#######################################  Definitions

exec 3> /dev/null # to avoid harcoding /dev/null everywhere. For tools' stderr.
SELF="${0##*/}" # buitin $(basename $0)

#user input
BOTTLE=
REGFILE=
BACKUP=
RESTORE=
REFRESH=
LIST=
OVERWRITE=
SKIP=cpl,dll,exe,lnk,msi
PREFIX=
UNDO=

#parameters
VERBOSE=1
TEST=
DEBUG=
GUI=1

#constants
CLASSTAG="WineImportedExtension"
regex='\\.[0-9A-Za-z_+%~-]+'  # thanks to mawk, [:alnum:] cannot be used :(
xdg_data_home=${XDG_DATA_HOME:-"$HOME"/.local/share}
xdg_data_dirs=${XDG_DATA_DIRS:-/usr/local/share:/usr/share}
winebottlehome=${WINEBOTTLEHOME:-"$xdg_data_home"/wineprefixes}

#global variables
command=
appname=
icon=
prefix="${WINEPREFIX:-"$HOME/.wine"}"
executable=
output="REGEDIT4"$'\n'$'\n'
outputfile=
sysreg=
native_ext=
wine_ext=
wine_class=
classes=
long=
savedlocale=
prev_ext=
tempfile=
n_ext=0
n_ref=0
n_ovr=0
n_skp=0


#######################################  Common functions (from lib, lowercase)

fatal()
{
	local message="${1:-}"
	local errorcode="${2:-1}"

	[[ "$message" ]] && printf "%s\n" "$SELF: ${message/%[[:punct:]]/}" >&2
	exit $errorcode
}

xdg_data_cat()
(
	local path="$1"
	local folder
	local xdgpath=${xdg_data_home}:${xdg_data_dirs}

	IFS=:
	for folder in $xdgpath ; do
		[[ -f "${folder%/}/$path" ]] && cat "${folder%/}/$path" 2>&3
	done
)

crc16()
(
	export LC_ALL=C
	local string="${1:-}"
	local crc=0
	local c i j xor_poly

	for ((i=0; i<${#string}; i++)); do
		c=$(printf "%d" "'${string:i:1}")
		for ((j=0; j<8; c>>=1, j++)); do
			(( xor_poly = (c ^ crc) & 1 ))
			(( crc >>= 1 ))
			(( xor_poly )) && (( crc ^= 0xA001 ))
		done
	done

	printf "%04X\n" "$crc"
)


#######################################  Functions (this script only)

Usage ()
{
	local explicit="${1:-}"

	cat <<-_USAGE
	Usage: $SELF [OPTIONS]
	_USAGE

	if [[ -z "$explicit" ]] ; then
		cat <<-_USAGE
		Try '$SELF --help' for more information.
		_USAGE
		return
	fi

	cat <<-_USAGE

	Imports all native file extensions (globs in xdg mime database) into Wine
	registry, setting file handler to winebrowser. This allows Wine applications
	to launch files in default native Linux applications

	Options:
	 -h, --help     print this message and exit
	 -q, --quiet    suppress informative messages (only displays errors)
	 -t, --test     simulation test, do not change wine registry
	 -d, --debug    displays debug information on parsed options and arguments

	 -k, --backup FILE
	     Backup current Wine Registry (system.reg) to FILE (WineRegistry format)
	     It actually just copies the PREFIX/system.reg file. This is the backup
	     file that must be provided for --restore option

	 -c, --savechanges FILE
	     Generates a registry file with all the applied changes (Windows format)
	     This is the file that must be provided for --undo option

	 -l, --list FILE
	     Generates a tab-delimited file with current native extensions and mime-
	     types and its associated wine classes and file handlers (if any)

	 -p, --prefix PATH
	     Sets the wine prefix to work. If used, enviroment variable WINEPREFIX
	     is ignored. It is overwritten by --bottle

	 -b, --bottle NAME
	     Sets the wine bottle to work. A bottle is a wine prefix currently
	     located at $winebottlehome/NAME
	     If used, enviroment variable WINEPREFIX is ignored. Overrides --prefix

	 -o, --overwrite EXT,...
	     Import listed native extension(s) even if there is already a registered
	     file handler in wine. Original class name will not be changed. If an
	     extension is registered to a class that has several extensions (for
	     example jpegfile for jpg and jpeg), all associated extensions will be
	     affected. Listed extensions that do not exist in native environment
	     will be silently ignored. List must be comma-delimited and EXT must be
	     the extension without any dots or asterisks.Ex: --overwrite jpg,gif,txt

	 -s, --skip EXT,...
	     Do not import native extension even if there is no registered class or
	     file handler in wine. Opposite of -o. List must be comma-delimited, EXT
	     must be the extension without any dots or asterisks. Ex: --skip dll,cpl
	     Extensions '$SKIP' are always skipped.

	 -G, --no-gui
	     Instead of the default GUI handler, register an alternate, simpler
	     handler that does not display any dialogs.

	 -a, --all
	     Re-generate all relevant registry keys. Default mode is to generate
	     only new entries. However, existing entries not generated with
	     $SELF will not be overwritten, unless --overwrite is also used.
	     Useful only if associations were manually edited in registry or to
	     set existing associations with a different --no-gui option.

	 -r, --restore FILE
	     Restores a backup registry file previously created with --backup option
	     It will revert all changes made by \$SELF. Any other
	     changes made after backup was created will be lost too, naturally.
	     To undo only the changes made \$SELF, and keep other
	     changes intact, use the --undo option.
	     IMPORTANT: RESTORE WILL DELETE NATIVE MIME TYPES CREATED BY *ALL* WINE
	     PREFIXES! Running winecfg (or regedit or uninstaller) in the other
	     prefixes will re-create them. This is already done for current
	     prefix (the one specified by --bottle, --prefix or WINEPREFIX
	     environment variable)

	 -u, --undo FILE (currently not implemented)
	     Undo changes made by \$SELF using the change file created
	     with --savechanges option. Tries its best to preserve untouched all
	     other changes that may have been introduced in registry since last run.
	     To fully (and safely) restore the registry to its previous state, use
	     the --restore option
	     THIS IS A HIGHLY EXPERIMENTAL FEATURE! USE AT YOUR OWN RISK!


	Examples and suggested uses:
	$SELF --debug --test --bottle test --backup system.reg.bak --list ext.txt
	$SELF --debug --test --bottle test --savechanges import.reg \\
	      --overwrite gif,jpg,png,txt --skip cpl,dll,lnk,msi	
	env PATH=$PATH:/opt/cxoffice/bin WINEBOTTLEHOME=$HOME/.cxoffice \\
	WINEPREFIX=$HOME/.cxoffice $SELF --bottle TotalCmd


	Author: Rodrigo Silva (MestreLion) <linux@rodrigosilva.com>
	License: GPLv3 or later, at your choice. See <http://www.gnu.org/licenses/gpl>
	_USAGE
}

PrintFlags()
{
	cat <<-_PRINTFLAGS
	Parsed Options  =$GETOPT
	Debug           = $DEBUG
	Verbose         = $VERBOSE
	Simulation Test = $TEST
	Refresh entries = $REFRESH
	GUI handler     = $GUI
	Wine Bottle     = $BOTTLE
	Wine Prefix     = $prefix
	Changes file    = $REGFILE
	Backup file     = $BACKUP
	Restore file    = $RESTORE
	List file       = $LIST
	_PRINTFLAGS
}

SetupWinePrefix()
{
	[[ -f "$sysreg" ]] && return

	[[ "$VERBOSE" ]] && {
		echo "Creating wine prefix $prefix"
		echo "This may take a while..."
	}
	wine wineboot 2>&3 || fatal "Could not create wine prefix $prefix"
	while [[ ! -f "$sysreg" ]] ; do : ; done ; sleep 3
}

SetupExecutable()
{
	local winexec

	if [[ "$GUI" || "$RESTORE" ]]; then
		executable="$prefix/dosdevices/c:/windows/system32/winebrowser-gui.exe"
	else
		executable="$prefix/dosdevices/c:/windows/system32/winebrowser.exe"
	fi

	[[ "$RESTORE" || "$UNDO" ]] && return

	winexec="$(wine winepath -w "$executable" 2>&3)"
	command="${winexec//\\/\\\\} \\\"%1\\\""

	appname="${executable##*/}"
	appname="${appname%%\.*}"

	icon="$(crc16 "$winexec")_${appname}.0"  # pray index will always be 0

	[[ "$GUI" && ! "$TEST" ]] || return

	cat > "$executable" <<-'_WINEBROWSERGUI'
	#!/bin/bash
	file=$(wine winepath -u "$1")
	winfile=${1//\\/\\\\}
	if ! xdg-open "$file"; then
	    mimetype=$(mimetype --brief "$file")
	    description=$(mimetype --brief --describe "$file")
	    msg="<big>Could not open <b>$winfile</b></big>\n\n"
	    msg+="There is no application installed for <b>$description</b> files.\n"
	    msg+="Do you want to search for an application to open this file?"
	    if zenity --question --no-wrap --text "$msg"; then
	        dbus-send --session --type=method_call --print-reply   \
	            --dest=org.freedesktop.PackageKit                  \
	            /org/freedesktop/PackageKit                        \
	            org.freedesktop.PackageKit.Modify.InstallMimeTypes \
	            int32:0                                            \
	            array:string:"$mimetype"                           \
	            string:"" &&
	        xdg-open "$file"
	    fi
	fi
	_WINEBROWSERGUI
	[[ $? -eq 0 ]] || fatal "Could not create $executable"
	chmod +x "$executable"

	[[ "$VERBOSE" ]] && echo "GUI launcher created in $winexec"
}

RegistryAddExt()
{
	local ext="$1"
	local class="$2"
	local undo="$3"

	output+="[HKEY_CLASSES_ROOT\\.${ext}]"$'\n'
	output+="@=\"${class}\""$'\n\n'
	[[ "$undo" ]] && output+="\"${CLASSTAG}_undo\"=\"${undo}\""$'\n'
}

RegistryAddClass()
{
	local class="$1"
	local handler="$2"

	output+="[HKEY_CLASSES_ROOT\\${class}\shell\\open\\command]"$'\n'
	output+="@=\"${handler}\""$'\n'
	output+=$'\n'
}

# winemnubuilder inhibitor
RegistryAddAssoc()
{
	local ext="$1"
	local class="$2"
	local mimetype="$3"

	output+="[HKEY_CURRENT_USER\\Software\\Wine\\FileOpenAssociations\\.${ext}]"$'\n'
	output+="\"AppName\"=\"${appname}\""$'\n'
	output+="\"MimeType\"=\"${mimetype}\""$'\n'
	output+="\"OpenWithIcon\"=\"${icon}\""$'\n'
	output+="\"ProgID\"=\"${class}\""$'\n\n'
}

#Do NOT use in subshells!
CreateTempFile()
{
	# See if this was called before
	if [[ "$tempfile" ]]; then
		# "push" previous file to end of array
		tempfile+=( "$tempfile" )
	else
		# set the trap with swapped quotes for delayed array evaluation
		trap 'rm -f -- "${tempfile[@]}"' EXIT
	fi

	# set new file
	tempfile=$(mktemp) || fatal "could not create temporary file"
}

Restore()
{
	local backup="$RESTORE"

	[[ "$TEST" ]] && {
		[[ "$VERBOSE" ]] && echo "Test run, registry was not changed"
		return 0
	}

	cp -- "$backup" "$sysreg" || fatal "Could not restore wine's registry"

	# Force wine to update extensions and mime types
	rm -f  -- "$xdg_data_home"/mime/packages/x-wine-*
	update-mime-database "$xdg_data_home"/mime
	wine winemenubuilder -a -r 2>&3

	# delete winebrowser-gui.exe
	rm -f -- "$executable"

	[[ "$VERBOSE" ]] && echo "Registry successfully restored"

	return 0
}

Undo()
{
	echo "Undo feature not implemented yet. Use --restore instead"
	Usage
	return
}


#######################################  Command Line parsing

long="help,quiet,test,debug,no-gui,bottle:,prefix:,savechanges:,overwrite:,\
skip:,backup:,restore:,list:,undo:,all"
GETOPT=$(getopt --alternative --name="$SELF" --options="hqtdGsb:p:k:c:o:r:l:u:"\
                --longoptions="$long" -- "$@"
) || { Usage ; exit 1; } # error in getopt

eval set -- "$GETOPT"

while true ; do
	arg="$1"
	shift
	case "$arg" in
	-q|--quiet      ) VERBOSE=               ;; # Unset verbose flag
	-t|--test       ) TEST=1                 ;; # Set simulation test
	-d|--debug      ) DEBUG=1                ;; # Set debug mode
	-G|--no-gui     ) GUI=                   ;; # Unset GUI flag
	-a|--all        ) REFRESH=1              ;; # Set refresh flag
	-p|--prefix     ) PREFIX="$1"   ; shift  ;; # Set Wine prefix
	-b|--bottle     ) BOTTLE="$1"   ; shift  ;; # Set Wine Bottle
	-c|--savechanges) REGFILE="$1"  ; shift  ;; # Set Registry change file
	-k|--backup     ) BACKUP="$1"   ; shift  ;; # Set Registry backup file
	-l|--list       ) LIST="$1"     ; shift  ;; # Set list file
	-r|--restore    ) RESTORE="$1"  ; shift  ;; # Set restore file
	-u|--undo       ) UNDO="$1"     ; shift  ;; # Set undo file
	-o|--overwrite  ) OVERWRITE="$1"; shift  ;; # Set list of exts to overwrite
	-s|--skip       ) SKIP+=",$1"   ; shift  ;; # Set list of exts to skip
	-h|--help       ) Usage "HELP"  ; exit   ;; # Help requested
	--              ) break                  ;;
	*               ) Usage         ; exit 1 ;;
	esac
done


#######################################  Main Body

[[ "$PREFIX" ]] && prefix="$PREFIX"
[[ "$BOTTLE" ]] && prefix="$winebottlehome/$BOTTLE"
[[ "$DEBUG"  ]] && { PrintFlags ; exec 3>&2 ; }

export WINEPREFIX="$prefix"

sysreg="$prefix/system.reg"

SetupWinePrefix
SetupExecutable

[[ "$UNDO"    ]] && { Undo    ; exit ; }
[[ "$RESTORE" ]] && { Restore ; exit ; }


# LC_ALL=C speed up processing and allow proper join
# But should not be set for Wine commands, since locale is relevant
savedlocale="${LC_ALL:-}"
export LC_ALL=C

# Native extensions (freedesktop.org xdg mime database globs)
# Format: <NUMBER>:$mimetype:$glob
native_ext=$(
	xdg_data_cat "mime/globs" |
	awk -F: -v regex="$regex" '
		$2 ~ "^\\*" regex "$" && $1 !~ /^#/ {
			ext = substr($2,3)
			lext = tolower(ext)
			if (native[lext]=="") {
				native[lext] = $1
				print ext "\t" $1
			}
		}' |
	sort --ignore-case
) 2>&3

[[ "$VERBOSE" ]] && printf "%4d %s\n" $(wc -l <<< "$native_ext")               \
                           "native extensions found (A)"

[[ "$DEBUG" ]] && echo "$native_ext" > "1native.txt"

# Wine extensions that have associated classes (from registry)
# Format:
# [Software\\Classes\\.$ext]<BLANK><NUMBER>
# @="$class"
wine_ext=$( { awk -F= -v regex="$regex" '
                  tolower($0)~"^\\[software\\\\\\\\classes\\\\\\\\" regex "]" {
                      split($0,a,/[].]/)
                      ext=a[2]
                      next
                  }
                  /^\[/ { ext="" }
                  /^@="/ {
                      if(ext) {
                          gsub(/^@="|"$/,"")
                          print tolower(ext) "\t" $0
                          ext=""
                      }
                  }' "$sysreg"                                                 \
              || fatal "Could not read wine registry $sysreg"
            }                                                                  \
            |  sort --unique --ignore-case                                     \
) 2>&3

[[ "$VERBOSE" ]] && printf "%4d %s\n" $(wc -l <<< "$wine_ext")                 \
                           "wine extensions with classes found (B)"

[[ "$DEBUG" ]] && echo "$wine_ext" > "2wine.txt"

# Loop the combined list
# (All native extensions with their corresponding wine class and file handler)
[[ "$LIST" ]] && { printf "EXT\tMIMETYPE\tCLASS\tFILE HANDLER\n" > "$LIST"     \
                   || fatal "could not create list file $LIST" ; }
while IFS=$'\t' read -r -s ext mimetype class; do

	[[ "$LIST" ]] && {
		if [[ "${LIST##*.}" = "ods" ]]; then
			printf "\"%s\"\t\"%s\"\t\"%s\"\n"                      \
	               "$ext" "$mimetype" "$class" >> "$LIST"
		else
			printf "%s\t%s\t%s\n"                                  \
	               "$ext" "$mimetype" "$class" >> "$LIST"
		fi
	}

	# Skip blacklisted extensions
	[[ ",${SKIP}," = *",${ext},"* ]] && { (( n_skp++ )); continue; }

	addreg=0
	oldclass=""

	# New extension (no class)
	if [[ -z "$class" ]]; then
		(( n_ext++ ))
		class="${CLASSTAG}.${ext}"
		addreg=1

	# Existing class, already our own handler: only (re-)add if refresh is set
	elif [[ "$class" = "${CLASSTAG}.${ext}" ]]; then
		if [[ "$REFRESH" ]]; then
			(( n_ref++ ))
			addreg=1
		fi

	# Existing extension, not our class, but in overwrite list: backup and add
	elif [[ ",${OVERWRITE}," = *",${ext},"* ]]; then
		(( n_ovr++ ))
		oldclass=$class
		class="${CLASSTAG}.${ext}"
		addreg=1
	fi

	if (( addreg )); then
		RegistryAddExt   "$ext"   "$class"  "$oldclass"
		RegistryAddClass "$class" "$command"
		RegistryAddAssoc "$ext"   "$class" "$mimetype"
		output+=$'\n\n'
	fi

done < <(
  join -a1 -t$'\t' -i -o1.1,1.2,2.2 <(echo "$native_ext") <(echo "$wine_ext") |
  sort --ignore-case
)

[[ "$VERBOSE" ]] && {
	printf "%4d %s\n" $n_ext "new extensions to be added (A - B)"

	[[ "$n_ref" -gt 0 ]] && printf "%4d %s\n" $n_ref                       \
	  "existing extensions to be refreshed"

	[[ "$n_ovr" -gt 0 ]] && printf "%4d %s\n" $n_ovr                       \
	  "existing windows extensions to be overwritten"

	[[ "$n_skp" -gt 0 ]] && printf "%4d %s\n" $n_skp                       \
	  "extensions skipped"
}

# Announce if list was created
[[ "$LIST" && "$VERBOSE" ]] && printf "List file was generated to \"$LIST\"\n"

# Backup current registry
[[ "$BACKUP" ]] && {
	cp -- "$sysreg" "$BACKUP" || fatal "Could not create backup file $BACKUP"
	[[ "$VERBOSE" ]] && printf "Current Registry was backed up to \"$BACKUP\"\n"
}

# Dump the output
[[ "$REGFILE" ]] && {
	printf '%s' "$output" > "$REGFILE" || fatal "Could not create registry file $REGFILE"
	[[ "$VERBOSE" ]] && printf "Registry file was generated to \"$REGFILE\"\n"
}

# Create temp file for regedit, or use the user generated
if [[ "$REGFILE" ]]; then
	outputfile="$REGFILE"
else
	CreateTempFile; outputfile="$tempfile"
	printf '%s' "$output" > "$outputfile"
fi

# Merge the registry
if [[ "$TEST" ]]; then
	[[ "$VERBOSE" ]] && echo "Test run, registry was not changed"
else
	export LC_ALL="$savedlocale"
	wine regedit "$outputfile" || fatal "Registry changes could not be merged into wine"
	sleep 3 && wine winemenubuilder -a -r 2>&3 # force refresh
	[[ "$VERBOSE" ]] && echo "Registry changes successfully merged into wine"
fi

exit 0


# Sample outputs (possibly outdated):

#rodrigo@desktop ~ $ wine-import-extensions -vdt -b test -k sys.reg.bak -l ext.txt
#Parsed Options  = -v -d --test -b 'test' -k 'sys.reg.bak' --list 'ext.txt' --
#Debug           = 1
#Verbose         = 1
#Simulation Test = 1
#GUI handler     =
#Wine Bottle     = test
#Wine Prefix     = /home/rodrigo/.local/share/wineprefixes/test
#Registry file   =
#Backup file     = sys.reg.bak
#List file       = ext.txt
#Restore file    =
# 760 native extensions found (A)
#  29 wine extensions with classes found (B)
#  21 wine classes with file handlers found (C)
# 696 extensions lacking class will be added (A \ B)
# 705 classes lacking file handler will be added (A ∩ B) \ C
#  35 extensions were duplicate with multiple mime-types and were ignored
#   1 classes were duplicate and thus ignored
#     (different extensions pointing to a class that was already processed)
#List file was generated to "ext.txt"
#Current Registry was backed up to "sys.reg.bak"
#Test run, registry was not changed
#rodrigo@desktop ~ $ echo $?
#0

#	rodrigo@desktop ~ $ /usr/bin/time -f'%E' wine-import -vdg -b teste -r reg.txt;echo $?
#	Parsed Options  = -v -d --gui --bottle 'teste' --regfile 'regfile.txt' --
#	Debug           = 1
#	Verbose         = 1
#	Non-Interactive =
#	Simulation Test =
#	GUI handler     = 1
#	Wine Bottle     = teste
#	Wine Prefix     = /home/rodrigo/.local/share/wineprefixes/teste
#	Executable      = winebrowser-gui
#	Registry file   = regfile.txt
#	 718 native extensions found (A)
#	  29 wine extensions with classes found (B)
#	  19 wine classes with file handlers found (C)
#	 689 extensions lacking class will be added (A \ B)
#	 701 classes lacking file handlers will be added (A ∩ B) \ C
#	Registry file "regfile.txt" was generated
#	0:00.28
#	0