clonezilla 3.27.16-2 / usr / sbin / ocs-btsrv

#!/bin/bash
# License: GPL 
# Author: Steven Shiau <steven _at_ clonezilla org>
# Description: Program to start BT service for restoring
# Including 3 parts:
# Part 1: create metainfo file (.torrent)
# Part 2: start tracker
# Part 3: start bitorrent client program

#
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
# metainfo_creator and bt_client is default defined in drbl-ocs.conf, or command option -t|--bt-client.
# 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, --metainfo-creator  PROG  Use PROG 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: ctorrent, lftp, btdownloadheadless, aria2c"
    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
    # Set the variable cmd_make_metainfo
    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"
    elif [ "$metainfo_creator" = "transmission-create" ]; then
      cmd_make_metainfo="transmission-create -p -c For_Clonezilla -t http://${btsrv_ip}:$port_b/announce -o $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
    if [ ! -e $ocsroot_btzone/${idir}.torrent ]; then
      echo "Generate metainfo file $ocsroot_btzone/${idir}.torrent by:"
      echo "$cmd_make_metainfo"
      eval $cmd_make_metainfo
    else
      # Maybe server IP address is changed. If so, regenerate it.
      # E.g. ...:announce35:http://192.168.22.250:6969/announce7:...
      existing_ip="$(LC_ALL=C strings $ocsroot_btzone/${idir}.torrent | grep -Eo ":announce[[:digit:]]+:http://[[:digit:]\.:]+.*announce[[:digit:]]+:" | awk -F":" '{print $4}' | sed -r -e "s|\/\/||g")"
      if [ "${existing_ip}" != "${btsrv_ip}" ]; then
        [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
        echo "Regenerating metainfo due to existing tracker \"${existing_ip}\" in $ocsroot_btzone/${idir}.torrent is different from detected one:\"${btsrv_ip}\"..."
        [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	rm -f $ocsroot_btzone/${idir}.torrent
        echo "$cmd_make_metainfo"
        eval $cmd_make_metainfo
      fi
    fi

    # 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 &"
    elif [ "$bt_client" = "aria2c" ]; then
      # //NOTE// The working dir is the parent dir.
      cmd_start_seeding="aria2c -D --seed-ratio=0.0 --enable-dht6=false --enable-dht=false --bt-seed-unverified $aria2c_extra_opt -d $ocsroot_btzone/"$(dirname ${idir})" $ocsroot_btzone/${idir}.torrent"
    else
      # Check the option -M of ctorrent, if it only allows 20-1000, modify it as "-M 20" from "-M 4", for example.
      # DRBL patched ctorrent allows -M to be 4-1000, otherwise it must be within 20-1000.
      # //NOTE// Here the white space after echo is a must. To avoid something like: echo "-n" -> make it as: echo " -n".
      ct_M="$(echo " $ctorrent_extra_opt" | grep -Ewo -- "[^[:space:]]*-M[[:space:]]+[[:digit:]]+[[:space:]]+" | awk -F" " '{print $2}')"
      if [ -n "$ct_M" ]; then
        if [ -n "$(LC_ALL=C strings `which ctorrent` | grep -Ew -- "^-%c argument must be between 20 and 1000")" ]; then
          # Lower limit for -M option is 20.
          if [ "$ct_M" -lt 20 ]; then
            # Force to make it as 20
            [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
            echo "The argument assigned in ctorrent_extra_opt for -M is $ct_M, which is lower than 20, the value allowed in the version of original ctorrent. Force to make it as 20."
            [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
	    ctorrent_extra_opt="$(echo " $ctorrent_extra_opt" | sed -r -e "s|-M[[:space:]]+[[:digit:]]+[[:space:]]+|-M 20 |" | sed -r -e "s|^[[:space:]]*||g")"
	  fi
        fi
      fi
      cmd_start_seeding="ctorrent -f -d $ctorrent_extra_opt -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 $! if the command is run in daemon mode.
    # 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.
    # TODO: Improve this. If user put more space or // in the path, it will be converted by system and the matching by grep won't work.
    ct_pid="$(LC_ALL=C ps -www -C "$bt_client" -o pid,cmd | grep -Ew "$cmd_start_seeding" | awk -F" " '{print $1}')"
    ct_pid="$(echo $ct_pid)"  # Make it in a line.
    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/aria2c 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/aria2c if available..."
    pkill -9 ocs-bttrack
    pkill -9 ctorrent
    pkill -9 lftp
    pkill -9 btdownloadheadless
    pkill -9 aria2c
  fi
  echo "Done!"
} # end of stop_bt_srv
#
show_status() {
  echo $msg_delimiter_star_line
  ps -www -C "ocs-bttrack" -C "aria2c" -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|--metainfo-creator) 
           shift; 
           if [ -z "$(echo $1 |grep ^-.)" ]; then
             # skip the -xx option, in case 
	     metainfo_creator="$1"
             shift;
           fi
           [ -z "$btlog_dir" ] && USAGE && exit 1
           ;;
   -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="$*"

# Load deafult values if not specified.
[ -z "$metainfo_creator" ] && metainfo_creator="$metainfo_creator_def"
[ -z "$bt_client" ] && bt_client="$bt_client_def"

# 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