#!/bin/bash
# License: GPL 
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to start BT service for restoring

#
DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"
. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
. /etc/drbl/drbl-ocs.conf
. $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Load the config in ocs-live.conf. This is specially for Clonezilla live. It will overwrite some settings of /etc/drbl/drbl-ocs.conf, such as $DIA...
[ -e "/etc/ocs/ocs-live.conf" ] && . /etc/ocs/ocs-live.conf

# Settings
bt_client="ctorrent"
metainfo_creator="mkttorrent"
# Flag to brutally kill the related service
brutal_kill="no"
# BT_PORT_INIT, p_length_ctorrent, p_length_mktorrent, ocsroot_btzone, btlog_dir are loaded from drbl-ocs.conf.

# Functions
USAGE() {
    echo "$ocs - To start the bittorrent for image dir"
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] [start|stop|status] [IMAGE] [DEV1] [DEV1]..."
    echo "Options:"
    echo "-c, --ctorrent-mkmetainfo    Use ctorrent instead of mktorrent to create the metainfo (.torrent) file"
    echo "-g, --log-dir DIR            Use DIR as the log dir. Default value is \"$btlog_dir\""
    echo "-m, --max-client-no NO       Assign the client number as NO. When this number of clients finishes downloading, BT tracker will exit so that client will stop, too."
    echo "-p, --bt-srv-ip  IP          Assign the IP address for using in tracker's URL. If not assigned, the IP address detected on the machine will be used"
    echo "-r, --bt-root DIR            Use DIR as the working dir. It's preceded to the files downloaded. Default value is \"$ocsroot_btzone\""
    echo "-t, --bt-client PROG         Use PROG instead of ctorrent as the bittorrent client program. Available PROG: lftp, btdownloadheadless"
    echo "IMAGE is the image name under \"$ocsroot_btzone\""
    echo "DEV* is the device name"
    echo "Ex:"
    echo "To start the bittorrent service for the image dir \"/home/partimag/btzone/myimg/sda1/\" and \"/home/partimag/btzone/myimg/sda5/\", assign the IP address 192.168.120.183 for the tracker's URL, run"
    echo "  $ocs -p 192.168.120.183 start myimg sda1 sda5"
    echo "To stop the bittorrent service for the image dir \"/home/partimag/btzone/myimg/sda5/\", run"
    echo "  $ocs stop my sda5"
    echo
} # end of USAGE
#
get_available_bt_port() {
  while nc -w 1 -z localhost $port_b; do
    # It's used. Use the next one by adding 2.
    port_b="$((port_b+2))"
  done
} # end of get_available_bt_port
#
wait_bt_port_open() {
  local port_query="$1"
  echo -n "Waiting for port $port_query of tracker to be started successfully..."
  while ! nc -w 1 -z localhost $port_query; do
    sleep 0.5
  done
  echo " done!"
} # end of wait_bt_port_open
#
get_idir_name() {
   local t_dir="$1"
   # Return global variable "idir_name"
   # t_dir is like: xenial-x64-20161104/sda1 xenial-x64-20161104/sda5
   t_dir_imgname="$(dirname $t_dir)"
   t_dir_devname="$(basename $t_dir)"
   if [ "$t_dir_imgname" = "." ]; then
     t_dir_imgname=""
   fi
   idir_name="${t_dir_imgname}~${t_dir_devname}"  # make it like: xenial-x64-20161104~sda1
} # end of get_idir_name
#
start_bt_srv() {
  local ipt
  if [ -z "$img_dir" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "No image dir was assigned!"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop!"
    exit 1
  fi
  # Check if the bt image dir exists
  for ipt in $parts; do
    if [ ! -d "$ocsroot_btzone/$img_dir/$(to_filename $ipt)" ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "Image dir \"$ocsroot_btzone/$img_dir/$(to_filename $ipt)\" not found!"
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      echo "$msg_program_stop!"
      exit 1
    fi
  done
  # Stop previous bt service first.
  stop_bt_srv

  # Assign an initial port no.
  port_b="${BT_PORT_INIT}"

  if [ -z "$btsrv_ip" ]; then
    # Decide the BT server IP address if it's not assigned.
    btsrv_ip="$(get-all-nic-ip -b | awk -F" " '{print $1}')"
  fi
  echo "Using IP address \"$btsrv_ip\" for BT server..."

  for ipt in $parts; do
    # Prepare port no
    get_available_bt_port # Obtain port_b

    idir="$img_dir/$(to_filename $ipt)"
    # idir_name is the name for log usage
    idir_name="${img_dir}:$(to_filename $ipt)"
    # Part 1: create metainfo file (.torrent)
    echo $msg_delimiter_star_line
    get_idir_name $idir # Obtian $idir_name
    rm -fv $ocsroot_btzone/${idir}.torrent
    if [ "$metainfo_creator" = "ctorrent" ]; then
      cmd_make_metainfo="ctorrent -t -p -c For_Clonezilla -l ${p_length_ctorrent} -u http://${btsrv_ip}:$port_b/announce -s $ocsroot_btzone/${idir}.torrent $ocsroot_btzone/$idir"
    else
      cpu_no="$(LC_ALL=C grep -Ew "^processor[[:space:]]+" /proc/cpuinfo | wc -l)"
      cmd_make_metainfo="mktorrent -p -c For_Clonezilla -t ${cpu_no} -l ${p_length_mktorrent} -a http://${btsrv_ip}:$port_b/announce -o $ocsroot_btzone/${idir}.torrent $ocsroot_btzone/$idir"
    fi
    echo "Generate metainfo file $ocsroot_btzone/${idir}.torrent by:"
    echo "$cmd_make_metainfo"
    eval $cmd_make_metainfo

    # Part 2: start tracker
    echo $msg_delimiter_star_line
    d_file="$btlog_dir/${idir_name}-dfile"
    log_file="$btlog_dir/${idir_name}-tracker.log"
    tracker_pid_log="$btlog_dir/${idir_name}-tracker.pid"
    if [ -n "$btclient_no" ]; then
      btclient_no_opt="--tracker_max_completed $btclient_no"
    fi
    rm -f $d_file $log_file $tracker_pid_log
    cmd_start_tracker="ocs-bttrack --port $port_b --dfile $d_file $btclient_no_opt --logfile $log_file --nat_check 0 --scrape_allowed full --allowed_dir $(dirname $(readlink -f $ocsroot_btzone/${idir}.torrent)) &"
    # Squeeze multiple spaces to one only
    cmd_start_tracker="$(echo "$cmd_start_tracker" | sed -r -e "s|[[:space:]]+| |g")"
    echo "Start tracker by:"
    echo "$cmd_start_tracker"
    eval $cmd_start_tracker
    tracker_pid="$!"
    echo "$tracker_pid" > $tracker_pid_log
    wait_bt_port_open $port_b

    # Part 3: start bitorrent client program
    echo $msg_delimiter_star_line
    ct_pid_log="$btlog_dir/${idir_name}-torrent.pid"
    if [ "$bt_client" = "lftp" ]; then
      cmd_start_seeding="lftp -c torrent -O $(dirname $ocsroot_btzone/${idir}) $ocsroot_btzone/${idir}.torrent"
    elif [ "$bt_client" = "btdownloadheadless" ]; then
      cmd_start_seeding="btdownloadheadless --saveas $ocsroot_btzone/${idir} $ocsroot_btzone/${idir}.torrent &"
    else
      cmd_start_seeding="ctorrent -d -s $ocsroot_btzone/${idir} $ocsroot_btzone/${idir}.torrent"
    fi
    echo "Start seeding by:"
    echo "$cmd_start_seeding"
    eval $cmd_start_seeding
    # Fail to get pid in this case by checking the $!.
    # Ref: http://unix.stackexchange.com/questions/90244/bash-run-command-in-background-and-capture-pid
    #ct_pid="$!"
    # Since start-stop-daemon is not available for all distributions, we use general solution here.
    ct_pid="$(LC_ALL=C ps -www -C "$bt_client" -o pid,cmd | grep -Ew "$cmd_start_seeding" | awk -F" " '{print $1}')"
    echo "$ct_pid" > $ct_pid_log
    echo $msg_delimiter_star_line
  done
} # end of start_bt_srv
#
stop_bt_srv() {
  local id_file=""
  local all_pid rc kill_cmd
  case "$img_dir" in 
     "all"|"") id_file="$btlog_dir/*.pid";;
            *)
  	       for ipt in $parts; do
		 idir="$img_dir/$(to_filename $ipt)"
                 # idir_name is the name for log usage
		 idir_name="${img_dir}:$(to_filename $ipt)"
                 id_file="$btlog_dir/${idir_name}-tracker.pid $btlog_dir/${idir_name}-torrent.pid"
               done
	       ;;
  esac

  # All
  if [ "$brutal_kill" = "no" ]; then
    echo "Stop service: ocs-bttrack and ctorrent/lftp/btdownloadheadless if available..."
    for i in $id_file; do
      [ ! -e "$i" ] && continue
      all_pid="$all_pid $(cat $i)"
    done
    if [ -n "$all_pid" ]; then
      echo "Found pid..."
      kill_cmd="kill -9 $all_pid"
      echo "Running: $kill_cmd"
      eval $kill_cmd
      rc=$?
      if [ "$rc" -eq 0 ]; then
        rm -f $id_file
      fi
    else
      echo "No pid was found... Do nothing."
    fi
  else
    echo "Brutally kill service ocs-bttrack and ctorrent/lftp/btdownloadheadless if available..."
    pkill -9 ocs-bttrack
    pkill -9 ctorrent
    pkill -9 lftp
    pkill -9 btdownloadheadless
  fi
  echo "Done!"
} # end of stop_bt_srv
#
show_status() {
  echo $msg_delimiter_star_line
  ps -www -C "ocs-bttrack" -C "ctorrent" -C "lftp" -C "btdownloadheadless" -o pid,tname,cmd
  echo $msg_delimiter_star_line
} # end of show_status

################
##### MAIN #####
################

ocs_file="$0"
ocs=`basename $ocs_file`
#
while [ $# -gt 0 ]; do
 case "$1" in
   -b|--brutal-kill) brutal_kill="yes"
	   shift;;
   -c|--ctorrent-mkmetainfo) metainfo_creator="ctorrent"
	   shift;;
   -g|--log-dir)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             btlog_dir="$1"
             shift;
           fi
           [ -z "$btlog_dir" ] && USAGE && exit 1
           ;;
   -m|--max-client-no)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             btclient_no="$1"
             shift;
           fi
           [ -z "$btclient_no" ] && USAGE && exit 1
           ;;
   -p|--bt-srv-ip)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             btsrv_ip="$1"
             shift;
           fi
           [ -z "$btsrv_ip" ] && USAGE && exit 1
           ;;
   -r|--bt-root)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             ocsroot_btzone="$1"
             shift;
           fi
           [ -z "$ocsroot_btzone" ] && USAGE && exit 1
           ;;
   -t|--bt-client)
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
             bt_client="$1"
             shift;
           fi
           [ -z "$bt_client" ] && USAGE && exit 1
           ;;
   -*)     echo "${0}: ${1}: invalid option" >&2
           USAGE >& 2
           exit 2 ;;
   *)      break ;;
 esac
done

action="$1"
shift
img_dir="$1"
shift
parts="$*"

# Strip the leading /dev/ if it's assigned. Make it like sda1, sdb1, not /dev/sda1, /dev/sdb1.
parts="$(echo $parts | sed -r -e "s|/dev/||g")"  # partition name, e.g. sda1, sda2

#check_if_root
ask_and_load_lang_set

case "$action" in
   start)  start_bt_srv ;;
    stop)  stop_bt_srv;;
  status)  show_status;;
       *)  USAGE; exit 1;;
esac
