# A library of commonly used functions written in expect.
# Copyright Brad Henry <brad@samba.org> 2006
# Released under the GNU GPL version 3 or later.

# This function maps a drive letter to a share point.
proc map_share { remote_prompt share_drive sharepoint username domain password } {
	set default_err_str "Unknown error in function map_share"
	set err_str $default_err_str

	set cmd "net use $share_drive $sharepoint $password /USER:$username@$domain\r\n"
	send $cmd

	expect {
		"The command completed successfully." {
			expect_prompt $remote_prompt
			set err_str "OK"
		} \
		"The local device name is already in use." {
			expect_prompt $remote_prompt
			set err_str "The device name $share_drive is already in use"
		} \
		"The network name cannot be found." {
			expect_prompt $remote_prompt
			set err_str "Sharepoint $sharepoint could not be found"
		} \
		timeout {
			set err_str "Function map_share timed out while mapping $share_drive to $sharepoint"
		}
	}
	return $err_str
}

# This function unmaps a drive letter from a share point.
proc unmap_share { remote_prompt share_drive } {
	set default_err_str "Unknown error in function unmap_share"
	set err_str $default_err_str

	set cmd "net use $share_drive /DELETE\r\n"
	send $cmd

	expect {
		"was deleted successfully." {
			expect_prompt $remote_prompt
			set err_str "OK"
		} \
		"NET HELPMSG 2250" {
			expect_prompt $remote_prompt
			set err_str "The network connection could not be found while unmapping $share_drive"
		} \
		timeout {
			set err_str "Function unmap_share timed out while unmapping $share_drive"
		}
	}
	return $err_str
}

# This function uses xcopy to copy a text file from one location on the
# remote windows host to another.
proc xcopy_file { remote_prompt in_filename out_filename xcopy_options } {
	set default_err_str "Unknown error in function xcopy_file"
	set err_str $default_err_str

	set cmd "xcopy $in_filename $out_filename $xcopy_options\r\n"
	send $cmd

	expect {
		"(F = file, D = directory)? " {
			set cmd "F\r\n"
			send $cmd
			expect {
				"1 File(s) copied\r\n\r\n" {
					expect_prompt $remote_prompt
					set err_str "OK"
				} \
				"0 File(s) copied\r\n\r\n" {
					expect_prompt $remote_prompt
					set err_str $default_err_str
				} \
				timeout {
					set err_str "Function xcopy_file has timed out while copying $in_filename"
				}
			}
		} \
		"1 File(s) copied\r\n\r\n" {
			expect_prompt $remote_prompt
			set err_str "OK"
		} \
		"0 File(s) copied\r\n\r\n" {
			expect_prompt $remote_prompt
			set err_str $default_err_str
		} \
		timeout {
			set err_str "Function xcopy_file timed out while copying $in_filename"
		}
	}
	return $err_str
}

# This function creates a temporary file on the remote windows host.
# The file contents are populated by a recursive directory listing of 
# the windows %HOMEDRIVE%.
proc create_tmp_file { remote_prompt filename } {
	set default_err_str "Unknown error in function create_tmp_file"
	set err_str $default_err_str

	set cmd "dir %HOMEDRIVE%\\ /S > $filename\r\n"
	send $cmd
	expect {
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			set err_str "Function create_tmp_file timed out while creating $filename"
		}
	}
	return $err_str
}

# This function compares two files on the remote windows host.
proc compare_files { remote_prompt file1 file2 } {
	set default_err_str "Unknown error in function compare_files"
	set err_str $default_err_str

	set cmd "fc $file1 $file2\r\n"
	send $cmd
	expect {
		"FC: no differences encountered\r\n\r\n\r\n" {
			expect_prompt $remote_prompt
			set err_str "OK"
		} \
		"\*\*\*\*\* $file1" {
			expect_prompt $remote_prompt
			set err_str "Files $file1 and $file2 differ"
		} \
		"\*\*\*\*\* $file2" {
			expect_prompt $remote_prompt
			set err_str "Files $file1 and $file2 differ"
		} \
		timeout {
			set err_str "Function compare_files timed out while comparing files $file1 and $file2"
		}
	}
	return $err_str
}

# This function deletes a file on the remote windows host.
proc delete_file { remote_prompt filename } {
	set default_err_str "Unknown error in function delete_file"
	set err_str $default_err_str

	set cmd "del $filename\r\n"
	send $cmd
	expect {
		"Could Not" {
			expect_prompt $remote_prompt
			set err_str $default_err_str
		} \
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			set err_str "Function delete_file timed oout while deleting $filename"
		}
	}
	return $err_str
}

# This function copies a text file over telnet from the local unix host
# to the remote windows host.
proc copy_file { remote_prompt in_filename out_filename } {
	set default_err_str "Unknown error in function copy_file"
	set err_str $default_err_str

	# The octal ASCII code for Control-Z is 032.
	set CTRLZ \032

	# Open local file and read contents.
	set in_file [open $in_filename r]
	set in_data [read $in_file]

	# Initiate copy on remote host.
	set cmd "copy con $out_filename\r\n"
	send $cmd

	# Separate $in_data into lines and send to remote host.
	set out_data [split $in_data "\n"]
	foreach out_line $out_data {
		send $out_line
		# We might as well do a unix -> windows line conversion.
		send "\r\n"
		# Are we overwriting an existing file?
		# If so, exit so we can handle it.
		expect {
			"(Yes/No/All)" {
				send "NO\r\n"
				expect_prompt $remote_prompt
				set err_str "File exists"
			} \
			$out_line {
				set err_str "OK"
			} \
			timeout {
				set err_str "Function copy_file timed out while copying $in_filename"
			}
		}
		if { $err_str != "OK" } {
			return $err_str
		} else {
			set err_str $default_err_str
		}
	}

	# ^Z\r to complete the transfer.
	send $CTRLZ
	send "\r"
	expect {
		"file(s) copied." {
			set err_str [expect_prompt $remote_prompt]
		} \
		$remote_prompt {
			set err_str $default_err_str
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Function copy_file timed out while finishing copy of $in_filename"
		}
	}
	return $err_str
}

# This function waits for the command prompt and reports an error on
# timeout.
proc expect_prompt { remote_prompt } {
	set default_err_str "Unknown error occurred while waiting for the command prompt"
	set err_str $default_err_str

	expect {
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			set err_str "Timeout occurred while waiting for the command prompt"
		}
	}
	return $err_str
}

# This function will create a telnet login shell to $remote_host as $username.
# If expected dialogue is not recieved, return with a specific error if one
# is recognized. Otherwise return a generic error indicating the function
# name.
proc telnet_login { remote_prompt remote_host username password } {

	set default_err_str "Unknown error in function telnet_login"
	set err_str $default_err_str

	set cmd "telnet $remote_host\r"
	send $cmd
	expect {
		"login: " {
			set err_str "OK"
		} \
		"Connection refused" {
			set err_str "Connection refused"
		} \
		"No route to host" {
			set err_str "No route to host"
		} \
		timeout {
			set err_str "Function telnet_login timed out while waiting for the login prompt"
		}
	}
	if { $err_str != "OK" } {
		# Return because something unexpected happened.
		return $err_str
	} else {
		# Reset err_str
		set err_str $default_err_str
	}

	set cmd "$username\r"
	send $cmd
	expect {
		"password: " {
			set err_str "OK"
		} \
		timeout {
			set err_str "Function telnet_login timed out while waiting for the password prompt"
		}
	}
	if { $err_str != "OK" } {
		return $err_str
	} else {
		set err_str $default_err_str
	}

	set cmd "$password\r"
	send $cmd
	expect {
		$remote_prompt {
			set err_str "OK"
		} \
		"Login Failed" {
			set err_str "Telnet login failed"
		} \
		timeout {
			set err_str "Function telnet_login timed out while waiting for the command prompt"
		}
	}
	return $err_str
}

proc create_directory { remote_prompt sharepath } {

	set default_err_str "Unknown error in function create_directory"
	set err_str $default_err_str

	set cmd "mkdir $sharepath\r\n"
	send $cmd
	expect  {
		"already exists" {
			expect_prompt $remote_prompt
			set err_str "Directory already exists"
		} \
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Timeout reached starting create_directory."
		}
	}
	return $err_str
}

proc delete_directory { remote_prompt sharepath } {

	set default_err_str "Unknown error in function delete_directory"
	set err_str $default_err_str

	set cmd "rmdir /S /Q $sharepath\r\n"
	send $cmd
	expect {
		"Access is denied." {
			expect_prompt $remote_prompt
			set err_str "Directory access is denied"
		} \
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Timeout reached in delete_directory"
		}
	}
	return $err_str
}

proc create_share { remote_prompt username sharepath sharename } {

	set default_err_str "Unknown error in function create_share"
	set err_str $default_err_str

	set cmd "net share $sharename=$sharepath /GRANT:$username,FULL\r\n"
	send $cmd
	expect {
		"was shared successfully." {
			set err_str [expect_prompt $remote_prompt]
		} \
		"NET HELPMSG 2118." {
			expect_prompt $remote_prompt
			set err_str "The name has already been shared"
		} \
		$remote_prompt {
			set err_str $default_err_str
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Timeout reached in create_share"
		}
	}
	return $err_str
}

proc delete_share { remote_prompt sharename } {

	set default_err_str "Unknown error in function delete_share"
	set err_str $default_err_str

	set cmd "net share $sharename /DELETE\r\n"
	send $cmd
	expect {
		"was deleted successfully." {
			set err_str [expect_prompt $remote_prompt]
		} \
		"does not exist." {
			expect_prompt $remote_prompt
			set err_str "The share does not exist"
		} \
		$remote_prompt {
			set err_str $default_err_str
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Timeout reached in delete_share"
		}
	}
	return $err_str
}

proc delete_hosts_entry { remote_prompt hosts_file_path backup_hosts_filename } {

	set default_err_str "Unknown error in function delete_hosts_entry"
	set err_str $default_err_str

	set cmd "cd $hosts_file_path\r\n"
	send $cmd
	expect {
		"." {
			expect_prompt $remote_prompt
			set err_str $default_err_str
		} \
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Timeout reached in delete_hosts_entry"
		}
	}
	if { $err_str != "OK" } {
		return $err_str
	} else {
		set err_str $default_err_str
	}

	set cmd "move /Y $backup_hosts_filename hosts\r\n"
	send $cmd
	expect {
		"1 file(s) moved." {
			set err_str [expect_prompt $remote_prompt]
		} \
		"cannot find the file specified." {
			expect_prompt $remote_prompt
			set err_str "File not found"
		} \
		$remote_prompt {
			set err_str $default_err_str
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Function delete_hosts_entry timed out while renaming $backup_hosts_filename"
		}
	}
	return $err_str
}

proc create_hosts_entry { remote_prompt hosts_file_path hostname ip \
				backup_hosts_filename } {

	set default_err_str "Unknown error in function create_hosts_entry"
	set err_str $default_err_str

	set cmd "cd $hosts_file_path\r\n"
	send $cmd
	expect	{
		"." {
			expect_prompt $remote_prompt
			set err_str $default_err_str
		} \
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Timeout reached in create_hosts_entry"
		}
	}
	if { $err_str != "OK" } {
		return $err_str
	} else {
		set err_str $default_err_str
	}

	set cmd "copy /Y hosts $backup_hosts_filename\r\n"
	send $cmd
	expect	{
		"1 file(s) copied." {
			set err_str [expect_prompt $remote_prompt]
		} \
		"cannot find the file specified." {
			expect_prompt $remote_prompt
			set err_str "File not found."
		} \
		$remote_prompt {
			set err_str $default_err_str
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Function create_hosts_entry timed out while copying hosts file"
		}
	}
	if { $err_str != "OK" } {
		return $err_str
	} else {
		set err_str $default_err_str
	}

	set cmd "echo $ip     $hostname     #smbtorture host. >> hosts\r\n"
	send $cmd
	expect	{
		$remote_prompt {
			set err_str "OK"
		} \
		timeout {
			expect_prompt $remote_prompt
			set err_str "Function create_hosts timed out while updating hosts file"
		}
	}
	return $err_str
}