# Andriy Zhugayevych  azh@ukr.net
# SSH network protocol support via PuTTY client
# created 25.12.2010, modified - see version below


SSH:=module() option package;

export ModuleLoad, Setup,
       opentunnel,  # port forwarding
       login, run, runftp,  # run interactive, command, and ftp SSH-clients
       pwd, dir, put, get, sget, # active remote directory, list remote directory, upload and download files
       FileExists, FileStatus,  # some commands
       qsub, qdel, qstat, pbsnodes, pbsreports;  # PBS routines

local version, requiredsetupversion,
      int, com, ftp, # filenames for interactive, command, and ftp SSH-clients
      md5, untar, ungz, # filenames for md5 checksum, untar and ungz programs
      sint, scom, sftp,
      tmp,  # filename for writing temporary Windows scripts

      host, # name@host
      ppk,  # private key file
      port,
      pbs,  # job scheduler (TORQUE, MOAB, SLURM)
      pbsopt, # PBS options
      cname,  # connection name (profile)
      ctable; # table of predefined connections
      


ModuleLoad:=proc()
local f;
version:=20220715;
requiredsetupversion:=20220715;
# Default settings
cname:="";
ctable:=table();
f:=cat(UserConfigurationFolder,"SSH_Setup.ini");
if FileTools[Exists](f) then read f else WARNING("SSH_Setup.ini is not found in UserConfigurationFolder=%1",UserConfigurationFolder) end;
NULL
end;


Setup:=proc(cname0::string:="",clist::list(string=[string,string,string]):=[],{setupversion::integer:=requiredsetupversion,printout::boolean:=false})
local sq,opts,auth;
if (setupversion<>requiredsetupversion) then WARNING("Setup versions do not match: submitted version %1, required version %2. Update SSH_Setup.ini",setupversion,requiredsetupversion) end;
if (clist<>[]) then ctable:=table(clist) end;
if (cname0<>"") then if assigned('ctable[cname0]') then cname,host,ppk,pbs:=cname0,op(ctable[cname0]) else error("Unrecognized cname: %1. Known cnames are %2",cname0,[indices(ctable,'nolist')]) end end;
sq:=ProcessSetupArgs([_rest],['int','com','ftp','md5','untar','ungz','tmp','host','port','ppk','pbs','pbsopt','cname'],['version','jobs']);
host,port:=op(..2,StringTools[Split](cat(host,":"),":"));
pbs,pbsopt:=op(..2,StringTools[Split](cat(pbs,":"),":"));
if (ppk="") then auth:=""
 elif FileTools[Exists](ppk) then auth:=cat(" -i ",ppk)
 else  auth:=cat(" -pw ",ppk);  if (length(ppk)>4 and ppk[-4..]=".ppk") then WARNING("Private key %1 does not exist. Provide the correct path to private key. Will be interpreted as password.",ppk) end end;
opts:=`if`(port="",NULL,cat(" -P ",port));
sint:=cat(int," -ssh",auth,opts," ",host);
scom:=cat(com," -ssh",auth,opts," ",host," ");
sftp:=cat(ftp,auth," -b ",tmp,opts," ",host);
if printout then
 printf("int=%s\ncom=%s\nftp=%s\nmd5=%s\nuntar=%s\nungz=%s\ntmp=%s\n",int,com,ftp,md5,untar,ungz,tmp);
 printf("host=%s, ppk=%s, port=%s, pbs=%s, pbsopt=%s\n",host,ppk,port,pbs,pbsopt);
 printf("cnames=[%{c,}s]\n",Vector([indices(ctable,'nolist')])) end;
sq
end;


opentunnel:=proc(destination::string,tunnel::string,tunnelppk::string,localport::posint:=`if`(type(port,posint),port,parse(port)),opts::string:=" -N",{printout::boolean:=false},$)
local s;
s:=cat(int," -ssh -i ",tunnelppk," -L ",localport,":",destination,opts," ",tunnel);
if printout then printf("%s\nTo close the tunnel close the opened window/process\n",s) end;
system[launch](s)
end;


login:=proc($) printf("Type \"exit\" to logout\n"); system[launch](sint) end;


run:=proc(command::string,$)
local ans;
ans:=ssystem(cat(scom,command));
if (ans[1]=0) then ans[2] else error "%1",ans end
end:


runftp:=proc(commands::list(string),$)
local ans;
WriteLines(tmp,[op(commands),"quit"],'overwrite');
if (host="a.zhugayevych@10.30.16.111") then ans:=system(sftp) else ans:=ssystem(sftp) end;
if (ans[1]=0) then ans[2] else error "%1",ans end
end:


pwd:=proc($)
local s,i;
s:=runftp(["pwd"]);
i:=SearchText("\nRemote directory is ",s);
`if`(i=0,s,s[i+21..-1])
end:


dir:=proc(wildcards::string:="",$)
local s,i;
s:=runftp([`if`(wildcards="","dir",cat("dir ",wildcards))]);
i:=SearchText("\n",s);
`if`(i=0,s,s[i+1..-1]) 
end:


put:=proc(files::string,cd::string:="",{lcd::string:=""},$)
runftp([`if`(cd="",NULL,cat("cd ",cd)),`if`(lcd="",NULL,cat("lcd ",lcd)),cat("mput ",files)])
end:


get:=proc(files::string,cd::string:="",{lcd::string:=""},$)
runftp([`if`(cd="",NULL,cat("cd ",cd)),`if`(lcd="",NULL,cat("lcd ",lcd)),cat("mget ",files)])
end:


sget:=proc(file::string,{lcd::string:="",compress::boolean:=false,unix2dos::boolean:=false,overwrite::boolean:=false},$)
local p2,s,p1,fn,fnx,cd,f2,s1,s2;
p2:=`if`(lcd="" or lcd[-1]="/" or lcd[-1]="\\",lcd,cat(lcd,"/"));
s:=FileStatus(file);
if (s=0) then error("Source file does not exist, %1",file) end;
p1,fn,fnx:=ExpandPath(file,"p,n,nx");
if (SearchText("*",file)>0) then  s:=8;  fn:=Substitute(fn,"*","_tmp_") end;
if unix2dos then run(cat("unix2dos ",file)) end;
if (compress or s=1 or s=8) then
 if FileTools[Exists](cat(p2,fnx)) then error("Target file exists (overwrite option is ignored for compressed transfer), %1%2",p2,fnx) end;
 if (s=8 and FileTools[ListDirectory](p2,'returnonly'=fnx)<>[]) then WARNING("Target files exist for mask %1%2, conflicts will be resolved according to untar=%3",p2,fnx,untar) end;
 run(cat("cd ",p1," && tar -zcf ",fn,".tar.gz ",fnx));
 sget(cat(p1,fn,".tar.gz"),':-lcd'=p2);
 run(cat("rm ",p1,fn,".tar.gz"));
 if (p2<>"") then  cd:=currentdir();  currentdir(p2) end;
 s2:=ssystem(cat(ungz ," ",fn,".tar.gz"));  fremove(cat(fn,".tar.gz"));
 s2:=ssystem(cat(untar," ",fn,".tar"   ));  fremove(cat(fn,".tar"   ));
 if (p2<>"") then currentdir(cd) end
else
 f2:=cat(p2,fnx);
 if (not(overwrite) and FileTools[Exists](f2)) then error("Target file exists, %1",f2) end;
 s1:=run(cat("md5sum ",file));
 s1:=StringTools[Split](s1," ")[1];
 runftp([`if`(lcd="",NULL,cat("lcd ",lcd)),cat("get ",file)]);
 s2:=ssystem(cat(md5," ",f2));
 s2:=StringTools[Split](StringTools[Split](s2[2],"\n")[-1]," ")[1];
 if (s1<>s2) then  fremove(f2);  error("Checksum error: %1<>%2",s1,s2) end end;
NULL
end:


#FileExists:=proc(filename::string,$) try run(cat("ls ",filename)) catch: return false end; true end:
FileExists:=proc(filename::string,$) local s; try s:=run(cat("ls -F -d ",filename)) catch: return false end; evalb(s<>"") end:

FileStatus:=proc(filename::string,$) local s; try s:=run(cat("ls -F -d ",filename)) catch: return 0 end; `if`(s="",0,`if`(SearchText("\n",s)>0,8,`if`(s[-1]="/",1,`if`(s[-1]="*",2,3)))) end:


qsub:=proc( script::string, jobname::string, pbsopt2::string:=pbsopt, {
 nodes::{string,nonnegint}:=0, ppn::nonnegint:=0, queue::string:="",
 mem::{numeric,undefined}:=undefined, pmem::{numeric,undefined}:=undefined,
 disk::{numeric,undefined}:=undefined,
 time::numeric:=0, file::{numeric,undefined}:=undefined, other::string:="",
 after::string:="", test::boolean:=false, printout::boolean:=false},$)
local s;
if (pbs="TORQUE" or pbs="MOAB") then
 s:=cat(
 `if`(nodes=0 and ppn=0,"",cat("nodes=",`if`(nodes=0,1,nodes),`if`(ppn=0,NULL,cat(":ppn=",ppn)))), #,`if`(queue="",NULL,cat(":",queue))
 `if`(mem>=0,cat(",mem=",`if`(mem<1,cat(max(1,round(1000*mem)),"mb"),cat(round(mem),"gb"))),NULL),
 `if`(pmem>=0,cat(",pmem=",round(1000*pmem),"mb"),NULL),
 `if`(disk>=0,cat(",ddisk=",round(disk),"gb"),NULL),
 `if`(file>=0,cat(",file=",round(file),"gb"),NULL),
 `if`(time>0,cat(",walltime=",FormatTime2(time*86400,'nodays')),NULL),
 other);
 s:=cat("qsub ",script," -N ",jobname,   # msub does not work on Pardus
 `if`(s="",NULL,cat(" -l ",s)),
 `if`(queue="",NULL,cat(" -q ",queue)),
 after)
elif (pbs="SLURM") then
 s:=cat("sbatch -J ",jobname,
 `if`(nodes=0,NULL,cat(" -N ",nodes)),
 `if`(ppn=0,NULL,cat(`if`(pbsopt2="t"," --ntasks-per-node="," -c "),ppn)),
 `if`(mem>=0,cat(" --mem=",`if`(mem<1,cat(max(1,round(1000*mem)),"mb"),cat(round(mem),"gb"))),NULL),
 `if`(pmem>=0,cat(" --mem-per-cpu=",round(1000*pmem),"mb"),NULL),
 `if`(time=0,NULL,cat(" -t ",FormatTime2(round(86400*time),'nodays'))),
 `if`(queue="",NULL,cat(" -p ",queue)),
 after," ",script)
else error("Unrecognized PBS: %1",pbs) end;
if (test or printout) then printf("%s\n",s) end;
if test then s else run(s) end
end:

qdel:=proc(jobnum::integer:=0,n::posint:=1,$)
local s,lsi,i;
if (pbs="TORQUE" or pbs="MOAB") then s:="qdel " elif (pbs="SLURM") then s:="scancel " else error("Unrecognized PBS: %1",pbs) end;
if (jobnum=0) then
 lsi:=sort(map2(op,1,select(v->v[2]="Q" or v[2]="PD",qstat('donotprint','fields'=["Job_Id","job_state"]))))
elif (jobnum<0) then
 lsi:=sort(map2(op,1,select(v->v[3][..3]="00:" and v[2]<>"C" and v[2]<>"E",qstat('donotprint','fields'=["Job_Id","job_state","resources_used"]))));
 lsi:=lsi[max(jobnum,-nops(lsi))..min(max(jobnum-n+1,-nops(lsi)),-1)]
else
 lsi:=[$jobnum..jobnum+n-1] end;
for i in lsi do  run(cat(s,i));  printf("deleted %d\n",i) end;
NULL
end:


qstat:=proc( sortbycols::list:=[3,1], { user::{string,undefined}:=undefined, donotprint::boolean:=false, sep::nonegint:=3,
  fields::list:=`if`(pbs="TORQUE",["Job_Id","Job_Name","Job_Owner","resources_used","job_state","queue","exec_host"],
                `if`(pbs="MOAB",["JobID","JobName","User","ReqProcs","State","Class","MasterHost","AWDuration"],[])),
   procfield::procedure:=proc(f::string,v,$)  local e;
   if type(v,string) then
    if   (f="Job_Id") then StringTools[Split](v,".")[1]
    elif (f="Job_Owner") then StringTools[Split](v,"@")[1]
    elif (f="resources_used") then `if`(v="",v,BasicTools[GetXMLChild](v,"cput"))
    elif (f="exec_host") then  e:=StringTools[Split](v,"+");  cat(StringTools[Split](e[1],"/")[1],"/",nops(e))
    elif (f="State") then piecewise(v="Running","R",v="Idle","Q",v)
    elif (f="AWDuration") then  e:=parse(v);  BasicTools[FormatTime2](e)
    else v end
   else "" end end },$)
local i,usr,s,xml,ls,lsw,e,fields0,M,n,maxlen, x,f,u,v,j;
if (user=undefined) then
  i:=SearchText("@",host);
  usr:=`if`(i>0,host[..i-1],"")
 else
  usr:=user end;
if (pbs="TORQUE") then
 s:=run("qstat -x"):
 if (s="") then return [] end;
 xml:=XMLTools[ParseString](s);
 ls:=GetXMLChild(xml,"Job",`if`(usr="",NULL,x->evalb(StringTools[Split](GetXMLChild(x,"Job_Owner"),"@")[1]=usr)),'all'):
 ls:=SortM([seq([seq(procfield(f,GetXMLChild(x,f)),f=fields)],x=ls)],sortbycols);
 if not(donotprint) then
  lsw:=[seq(max(seq(length(e[i]),e=ls))+sep,i=1..nops(fields))];
  for e in ls do printf(cat(seq(sprintf("%-*s",lsw[i],e[i]),i=1..nops(fields)),"\n")) end end;
 ls
elif (pbs="MOAB") then
 s:=run(`if`(usr="","showq --xml",cat("showq --xml -w user=",usr)));
 xml:=XMLTools[ParseString](s);
 ls:=remove(`=`,GetXMLChild(xml,"queue",'all'),"");
 ls:=[seq(seq(XMLTools[AttributeTable](v),v=XMLTools[GetChildByName](u,"job")),u=ls)];
 ls:=SortM(map(v->[seq(procfield(f,v[f]),f=fields)],ls),sortbycols);
 if not(donotprint) then
  lsw:=[seq(max(seq(length(e[i]),e=ls))+sep,i=1..nops(fields))];
  for e in ls do printf(cat(seq(sprintf("%-*s",lsw[i],e[i]),i=1..nops(fields)),"\n")) end end;
 ls
elif (pbs="SLURM") then
 if (usr="") then
  s:=run("squeue");
  if not(donotprint) then printf("%s\n",s) end;
  ls:=StringTools[Split](s,"\n")[2..];
  ls:=[seq(map(Trim,[s[1..7],s[18..26],s[27..35],s[39..49],s[36..38],s[8..17],s[62..]]),s=ls)];
  fields0:=["Job_Id","Job_Name","Job_Owner","resources_used","job_state","queue","exec_host"];
  M:=map(v->ListTools[Search](v,fields0),fields);
  if has(M,0) then error("Recognized fields are %1 but received %2",fields0,fields) end;
  map(v->v[M],ls)
 else
  s:=run(cat("squeue -u ",usr," -o '%i %j %M %t %D %C %P %R'"));
  ls:=StringTools[Split](s,"\n")[2..];
  ls:=map(StringTools[Split],ls," ");
  if not(donotprint or ls=[]) then
   n:=min(map(nops,ls));
   maxlen:=[seq(max(seq(length(ls[j][i]),j=1..nops(ls))),i=1..n)];
   for e in ls do  for i from 1 to n do printf(`if`(i=3," %*s "," %-*s "),maxlen[i],e[i]) end;  printf("\n") end end;
  ls end
else error("Unrecognized PBS: %1",pbs) end
end:


pbsnodes:=proc({ donotprint::boolean:=false, maxnp::posint:=64, allnodes::boolean:=false, nnp::string:="", # nnp = node name prefix
  fields::list:=["name","state","np","jobs"],
  procfield::procedure:=proc(f::string,v:="",$) if (f="jobs") then map(StringTools[Trim],StringTools[Split](v,",")) else v end end },$)
local getcores,s,xml,ls,nodes,i1,i2,i3,i4,n,minlen,maxlen,template,i0,i,nd,nb,nf,freeprocstat,node,np,lsb,nfp,s1,s2, x,f;
getcores:=proc(s0::string,$) local s,v;
 s:=StringTools[Split](s0,"/")[1];
 v:=sscanf(s,"%d-%d");
 if (v=[] or v=0) then NULL elif (nops(v)=1) then v[1] else $v[1]..v[2] end end;
if (pbs="TORQUE" or pbs="MOAB") then
 s:=run("pbsnodes -x");
 xml:=XMLTools[ParseString](s);
 ls:=GetXMLChild(xml,"Node",'all');
 nodes:=[seq([seq(procfield(f,GetXMLChild(x,f)),f=fields)],x=ls)];
 if not(allnodes) then nodes:=select(v->v[1][..length(nnp)]=nnp,nodes) end;
 if not(donotprint) then
  for i1 from 1 to nops(fields) while (fields[i1]<>"name" ) do end;  if (i1>nops(fields)) then error("No name in fields, %1",fields) end;
  for i2 from 1 to nops(fields) while (fields[i2]<>"state") do end;  if (i2>nops(fields)) then error("No state in fields, %1",fields) end;
  for i3 from 1 to nops(fields) while (fields[i3]<>"np"   ) do end;  if (i3>nops(fields)) then error("No np in fields, %1",fields) end;
  for i4 from 1 to nops(fields) while (fields[i4]<>"jobs" ) do end;  if (i4>nops(fields)) then error("No jobs in fields, %1",fields) end;
  n:=nops(nodes);
  minlen,maxlen:=min(seq(length(v[i1]),v=nodes)),max(seq(length(v[i1]),v=nodes));
  template:="";
  for i0 from 1 to minlen do
   s:=nodes[1][i1][i0];
   for i from 2 to n while (nodes[i][i1][i0]=s) do end;
   if (i<=n) then break else template:=cat(template,s) end end;
  template:=cat(template,"X"$(maxlen-i0+1));
  nd,nb,nf:=0,0,0;
  freeprocstat:=Array(0..maxnp,i->[]);
  for node in nodes do
   np:=op(sscanf(node[i3],"%d"));
   lsb:=sort(map(getcores,node[i4]));
   if (node[i2]="down" or node[i2]="offline") then
    if (lsb<>[]) then WARNING("Node %1 is %2 but jobs are running: %2",node[i1],node[i2],node[i4]) end;
    nd:=nd+np
   elif (node[i2]="job-exclusive") then
    if (nops(lsb)<>np) then WARNING("Node %1 is job-exclusive but not all cores are used: %2",node[i1],node[i4]) end;
    nb:=nb+np
   else
    nfp:=np-nops(lsb);
    nf:=nf+nfp;  nb:=nb+np-nfp;
    freeprocstat[nfp]:=[op(freeprocstat[nfp]),node[i1][i0..]] end end;
  s1,s2:="","";
  for nfp from 1 to maxnp do if (nops(freeprocstat[nfp])>0) then
   s1:=cat(s1,", (",nfp,")=",nops(freeprocstat[nfp]));
   s2:=cat(s2,", (",nfp,")=",StringTools[DeleteSpace](sprintf("%a",sort(freeprocstat[nfp])))) end end;
  printf("%d nodes, %d cores: %d down, %d busy, %d free\n",nops(nodes),nd+nb+nf,nd,nb,nf);
  printf("free-core statistics: (down)=%d, (busy)=%d, %s\n",nops(nodes)-add(nops(freeprocstat[nfp]),nfp=0..maxnp),nops(freeprocstat[0]),s1[2..]);
  printf("List of nodes (%s):%s\n",template,s2[2..]) end;
 Sort(nodes,e->op(sscanf(e[1],cat(nnp,"%d"))))
elif (pbs="SLURM") then
 s:=run("sinfo --Node");
 printf("%s\n",s)
else error("Unrecognized PBS: %1",pbs) end
end:

pbsreports:=proc({donotprint::boolean:=false},$)
local getjobnum,getjobname,ff,s,ls,v,d1;
getjobname:=fn->ExpandPath(fn,"n");  # `if`(pbs="TORQUE", ExpandPath(fn,"n"), `if`(pbs="SLURM", ExpandPath(SSH[run](cat("sed -n 3p ",fn)),"n") ,""));
getjobnum:=proc(fn::string,$)
 local v;
 if   (pbs="TORQUE" or pbs="MOAB" or pbs="SLURM") then v:=sscanf(StringTools[Split](fn,".")[-1],"e%d")
# elif (pbs="SLURM")  then v:=sscanf(StringTools[Split](fn,".")[-1],"%d")
 else 0 end;
 `if`(v=[],0,op(v))
 end;
ff:=proc(s::string)
 local ls,fn,jobnum,jobname,size;
 if s[1]="d" then return NULL end;
 ls:=remove(`=`,StringTools[Split](s," "),"");
 fn:=ls[-1];
 if (fn[1]="." or fn[1]="_") then return NULL end;
 jobnum:=getjobnum(fn);
 if (jobnum=0) then return NULL end;
 jobname:=getjobname(fn);
 if (jobname="") then jobname:=cat("--- Unrecognized job name in ",fn) end;
 size:=sscanf(ls[-5],"%d");
 if (size=[]) then error("Unrecognized filesize: %1",ls[-5]) else size:=op(size) end;
 [jobnum,size,jobname]
 end;
s:=dir();
ls:=StringTools[Split](s,"\n")[2..];
ls:=map(ff,ls);
ls:=SortM(select(v->v[1]>0,ls),[3,1]);
if not(donotprint or ls=[]) then
 d1:=1+floor(log10(max(seq(v[1],v=ls))));
 for v in ls do printf("%*d%4d%s %s\n",d1,v[1],op(`if`(v[2]<1000,[v[2]," "],`if`(v[2]<1000000,[round(v[2]/1000),"K"],[round(v[2]/1000000),"M"]))),v[3]) end end;
ls
end:



ModuleLoad()

end module: