/*## (c) SECURITY EXPLORATIONS    2018 poland                                #*/
/*##     http://www.security-explorations.com                                #*/

/* THIS SOFTWARE IS PROTECTED BY DOMESTIC AND INTERNATIONAL COPYRIGHT LAWS    */
/* UNAUTHORISED COPYING OF THIS SOFTWARE IN EITHER SOURCE OR BINARY FORM IS   */
/* EXPRESSLY FORBIDDEN. ANY USE, INCLUDING THE REPRODUCTION, MODIFICATION,    */
/* DISTRIBUTION, TRANSMISSION, RE-PUBLICATION, STORAGE OR DISPLAY OF ANY      */
/* PART OF THE SOFTWARE, FOR COMMERCIAL OR ANY OTHER PURPOSES REQUIRES A      */
/* VALID LICENSE FROM THE COPYRIGHT HOLDER.                                   */

/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS    */
/* OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,*/
/* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    */
/* SECURITY EXPLORATIONS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, */
/* WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF  */
/* OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE     */
/* SOFTWARE.                                                                  */

import java.lang.*;
import java.util.*;
import java.net.*;
import java.io.*;
import java.nio.*;
import java.nio.channels.*;

public class BoxServer {
 public static final int BUFFER_SIZE    = 2*Proto.BUFFER_SIZE;
 public static final int MAX_DESC_SIZE  = 0x40;

 public static final int HEADER_SIZE    = 0x06;
 public static final int MAX_DATA_SIZE  = BUFFER_SIZE;

 public static final int RS_READ_HEADER = 0x01;
 public static final int RS_READ_DATA   = 0x02;

 public static class Client {
  int type;
  String desc;
  int addr;

  SocketChannel ch;

  int read_state;
  int msg_size;
  int msg_type;

  byte hdr[];
  int hdr_size;

  ByteBuffer buf;
  int read_size;
  boolean has_msg;

  public Client(int type,int addr,String desc,SocketChannel ch) {
   this.type=type;
   this.addr=addr;
   this.desc=desc;
   this.ch=ch;

   read_state=RS_READ_HEADER;
   msg_type=-1;
   msg_size=-1;

   hdr=new byte[HEADER_SIZE];
   hdr_size=0;

   buf=ByteBuffer.allocate(BUFFER_SIZE);
   read_size=0;

   has_msg=false;
  }

  public String addr() {
   int b1=(addr>>24)&0xff;
   int b2=(addr>>16)&0xff;
   int b3=(addr>>8)&0xff;
   int b4=addr&0xff;

   String s=b1+"."+b2+"."+b3+"."+b4;

   return s;
  }
 }

 private static ServerSocketChannel srv_ch;
 private static Selector selector;

 private static ByteBuffer buf;
 private static ByteBuffer ecm_buf;

 private static boolean running;
 private static int req_id;

 private static Vector clients;
 private static Hashtable map;

 private static Client console;
 private static Client cur_box;

 private static Hashtable ecm_providers;//those providing ECM data (locator->Client)
 private static Hashtable ecm_receivers;//those willing to receive ECM data (locator->vector of Clients)

 private static Hashtable names_lookup;//channel names by locators lookup

 static {
  running=false;
  req_id=1;

  cur_box=console=null;

  clients=new Vector();
  map=new Hashtable();

  ecm_providers=new Hashtable();
  ecm_receivers=new Hashtable();

  names_lookup=new Hashtable();

  buf=ByteBuffer.allocate(BUFFER_SIZE);
  ecm_buf=ByteBuffer.allocate(BUFFER_SIZE);
 }

 private static Client find_client(SocketChannel ch) {
  return (Client)map.get(ch);
 }

 private static void add_client(Client cl) {
  clients.addElement(cl);
  map.put(cl.ch,cl);

  if (cl.type==Proto.CLT_CONSOLE) {
   console=cl;
  }
 }

 private static void del_client(Client cl) {
  if (find_client(cl.ch)!=null) {
   clients.removeElement(cl);
   map.remove(cl.ch);

   if (cl.type==Proto.CLT_CONSOLE) {
    console=null;
   }

   if (cur_box==cl) {
    cur_box=null;
   }
  }  
 }

 private static String lookup_service(String locator) {
  return (String)names_lookup.get(locator);
 }

 private static void cache_service(String locator,String name) {
  if (lookup_service(locator)==null) {
   names_lookup.put(locator,name);
  }
 }

 private static void add_provider(String locator,Client cl) {
  synchronized(ecm_providers) {
   ecm_providers.put(locator,cl);
  }
 }

 private static void del_provider(String locator) {
  synchronized(ecm_providers) {
   ecm_providers.remove(locator);
  }
 }

 private static Vector lookup_receivers(String locator) {
  return (Vector)ecm_receivers.get(locator);
 }

 private static boolean new_receiver(Vector v,Client recv) {
  boolean res=true;

  for(int i=0;i<v.size();i++) {
   Client cl=(Client)v.elementAt(i);
   if (cl==recv) {
    res=false;
    break;    
   }
  }

  return res;
 }

 private static void add_receiver(String locator,Client cl) {
  synchronized(ecm_receivers) {
   Vector v=(Vector)lookup_receivers(locator);

   if (v==null) {
    v=new Vector();
    ecm_receivers.put(locator,v);

    v.addElement(cl);
   } else
    
   if (new_receiver(v,cl)) {
    v.addElement(cl);
   }
  }
 }

 private static void del_receiver(String locator,Client cl) {
  synchronized(ecm_receivers) {
   Vector v=(Vector)lookup_receivers(locator);

   if (v!=null) {
    v.removeElement(cl);
   }
  }
 }

 private static boolean start_server(int port) {
  try {
   srv_ch=ServerSocketChannel.open();
   srv_ch.configureBlocking(false);

   srv_ch.socket().bind(new InetSocketAddress(port));

   selector=Selector.open();
   srv_ch.register(selector,SelectionKey.OP_ACCEPT);

   running=true;
  } catch(Throwable t) {}

  return running;
 }

 private static int read_client_type(WrappedInput wi) throws Throwable {
  int type=wi.get_byte();

  if ((type==Proto.CLT_BOX)||(type==Proto.CLT_CONSOLE)) return type;

  return -1;
 }

 private static String read_client_desc(WrappedInput wi) throws Throwable {
  return wi.read_string_with_size(MAX_DESC_SIZE);
 }

 private static void drop_channel(SocketChannel ch) throws Throwable {
  ch.close();
 }

 private static void drop_client(Client cl) throws Throwable {
  del_client(cl);
  drop_channel(cl.ch);

  System.out.println("client dropped: "+cl.type+" "+cl.desc);
 }

 private static boolean allowed_msg_type(int type) {
  boolean allowed=false;

  switch(type) {
   case Proto.CMD_CLIENTS_COUNT:
   case Proto.CMD_CLIENT_ID:
   case Proto.CMD_LIST:
   case Proto.CMD_GO:
   case Proto.CMD_EXIT:
   case Proto.CMD_ECMROUTES:
   case Proto.CMD_ANNOUNCEMENT:
   case Proto.CMD_DATA:
    allowed=true;
    break;
  }

  return allowed;
 }

 private static boolean box_msg_size(int size) {
  if (size<MAX_DATA_SIZE) return true;
   else return false;
 }

 private static void reset_msg(Client cl) {
  cl.read_state=RS_READ_HEADER;
  cl.msg_type=-1;
  cl.msg_size=-1;
    
  cl.hdr_size=0;
  cl.read_size=0;

  cl.buf.clear();
  cl.has_msg=false;
 }

 private static void write_fully(Client to,ByteBuffer buf) throws Throwable {
  int total=buf.remaining();

  while(total>0) {
   int cnt=to.ch.write(buf);
   total-=cnt;
  }
 }

 private static void discard_msg(Client cl) throws Throwable {
  cl.buf.position(cl.buf.limit());
 }

 private static void forward_msg(Client from,Client to) throws Throwable {
  if (to!=console) {
   ByteBuffer hdrbb=ByteBuffer.wrap(from.hdr);

   write_fully(to,hdrbb);
  }

  write_fully(to,from.buf);

  reset_msg(to);
 }

 private static void process_box_data_msg(Client cl) throws Throwable {
  if (console!=null) forward_msg(cl,console);
   else discard_msg(cl);
 }

 private static void process_new_service_announcement(Client cl,WrappedInput wi) throws Throwable {
  /* TODO*/
  throw null;
 }

 private static void send_ecm_message(Client cl,int stamp,int cw[]) { 
  try {
   ecm_buf.clear();

   ByteBufferOutput wo=new ByteBufferOutput(ecm_buf);
   wo.with_size();

   wo.set_data_type(Proto.CMD_ANNOUNCEMENT);

   wo.write_byte(Proto.AN_ECM_DATA);

   wo.write_int(stamp);

   for(int i=0;i<cw.length;i++) {
    wo.write_int(cw[i]);
   }

   wo.flush();

   write_fully(cl,wo.buf());
  } catch(Throwable t) {
   t.printStackTrace();
  }
 }

 private static void propagate_ecm(String locator,int stamp,int cw[]) throws Throwable {
  synchronized(ecm_receivers) {
   Vector v=lookup_receivers(locator);
  
   if (v!=null) {
    for(int i=0;i<v.size();i++) {
     Client cl=(Client)v.elementAt(i);
     send_ecm_message(cl,stamp,cw);    
    }
   }
  }
 }

 private static void process_dtv_ecm_announcement(Client cl,WrappedInput wi) throws Throwable {
  String locator=wi.read_string();
  int stamp=wi.read_int();

  int cw[]=new int[4];
  cw[0]=wi.read_int();
  cw[1]=wi.read_int();
  cw[2]=wi.read_int();
  cw[3]=wi.read_int();

  propagate_ecm(locator,stamp,cw);
 }

 private static void process_ecm_forwarder_start_announcement(Client cl,WrappedInput wi) throws Throwable {
  String srvname=wi.read_string();
  String locator=wi.read_string();

  cache_service(locator,srvname);
  add_provider(locator,cl);
 }

 private static void process_ecm_forwarder_stop_announcement(Client cl,WrappedInput wi) throws Throwable {
  String srvname=wi.read_string();
  String locator=wi.read_string();

  del_provider(locator);
 }

 private static void process_ecm_receiver_start_announcement(Client cl,WrappedInput wi) throws Throwable {
  String srvname=wi.read_string();
  String locator=wi.read_string();

  cache_service(locator,srvname);
  add_receiver(locator,cl);
 }

 private static void process_ecm_receiver_stop_announcement(Client cl,WrappedInput wi) throws Throwable {
  String srvname=wi.read_string();
  String locator=wi.read_string();

  del_receiver(locator,cl);
 }

 private static void process_announcement_msg(Client cl) throws Throwable {
  WrappedInput wi=new ByteBufferInput(cl.buf);
  int type=wi.get_byte();

  switch(type) {
   case Proto.AN_NEW_SERVICE:
    process_new_service_announcement(cl,wi);
    break;

   case Proto.AN_ECM_DATA:
    process_dtv_ecm_announcement(cl,wi);
    break;

   case Proto.AN_ECM_FORWARDER_STARTED:
    process_ecm_forwarder_start_announcement(cl,wi);
    break;

   case Proto.AN_ECM_FORWARDER_STOPPED:
    process_ecm_forwarder_stop_announcement(cl,wi);
    break;

   case Proto.AN_ECM_RECEIVER_STARTED:
    process_ecm_receiver_start_announcement(cl,wi);
    break;

   case Proto.AN_ECM_RECEIVER_STOPPED:
    process_ecm_receiver_stop_announcement(cl,wi);
    break;
   }
 }

 private static void process_box_msg(Client cl) throws Throwable {
  switch(cl.msg_type) {
   case Proto.CMD_ANNOUNCEMENT:
    process_announcement_msg(cl);
    break;
   case Proto.CMD_DATA:
    process_box_data_msg(cl);
    break;
  }
 }

 private static void process_clients_count_msg(Client cl) throws Throwable {
  ByteBufferOutput wo=new ByteBufferOutput(cl.buf);

  wo.write_int(clients.size()-1);

  wo.flush();

  write_fully(console,wo.buf());
 }

 private static void process_client_id_msg(Client cl) throws Throwable {
  int id=-1;

  for(int i=0;i<clients.size();i++) {
   Client c=(Client)clients.elementAt(i);

   if (c==cur_box) {
    id=i;
    break;
   }
  }

  ByteBufferOutput wo=new ByteBufferOutput(cl.buf);

  wo.write_int(id);

  wo.flush();

  write_fully(console,wo.buf());
 }

 private static void process_list_msg(Client cl) throws Throwable {
  ByteBufferOutput wo=new ByteBufferOutput(cl.buf);

  for(int i=0;i<clients.size();i++) {
   Client c=(Client)clients.elementAt(i);

   String line="";

   if (c==cur_box) {
    line+="*[";
   } else {
    line+=" [";
   }

   line+=ConsoleUtils.padded_int(i,2)+"] "+ConsoleUtils.padded_string(c.addr(),15)+" \""+c.desc+"\"";

   wo.write_string(line);
  }

  wo.write_int(-1);

  wo.flush();

  write_fully(console,wo.buf());
 }

 private static String channel_name(String locator) {
  return (String)names_lookup.get(locator);
 }

 private static void print_established_ecm_routes(ByteBufferOutput wo,Hashtable seen) throws Throwable {
  for(Enumeration e=ecm_providers.keys();e.hasMoreElements();) {
   String locator=(String)e.nextElement();
   Client fwd_cl=(Client)ecm_providers.get(locator);

   seen.put(locator,locator);
   
   String chname=channel_name(locator);
   boolean first=true;

   Vector v=(Vector)ecm_receivers.get(locator);

   if ((v==null)||(v.size()==0)) {
    String line=ConsoleUtils.padded_string(fwd_cl.addr(),15)+" \""+ConsoleUtils.padded_string(chname,20)+"\" -> ";
    wo.write_string(line);
    continue;
   }

   for(int i=0;i<v.size();i++) {
    Client rcv_cl=(Client)v.elementAt(i);

    String line="";

    if (first) {
     line+=ConsoleUtils.padded_string(fwd_cl.addr(),15)+" \""+ConsoleUtils.padded_string(chname,20)+"\" -> "+ConsoleUtils.padded_string(rcv_cl.addr(),15);
     first=false;
    } else {
     line+=ConsoleUtils.padded_string("",38)+" -> "+ConsoleUtils.padded_string(rcv_cl.addr(),15);
    }

    wo.write_string(line);
   }   
  }
 }

 private static void print_pending_ecm_routes(ByteBufferOutput wo,Hashtable seen) throws Throwable {
  for(Enumeration e=ecm_receivers.keys();e.hasMoreElements();) {
   String locator=(String)e.nextElement();

   if (seen.get(locator)==null) {
    Vector v=(Vector)ecm_receivers.get(locator);

    if ((v==null)||(v.size()==0)) continue;

    String chname=channel_name(locator);

    for(int i=0;i<v.size();i++) {
     Client rcv_cl=(Client)v.elementAt(i);
 
     String line=ConsoleUtils.padded_string("",15)+" \""+ConsoleUtils.padded_string(chname,20)+"\" -> "+ConsoleUtils.padded_string(rcv_cl.addr(),15);
     wo.write_string(line);
    }
   }
  }
 }

 private static void process_ecmroutes_msg(Client cl) throws Throwable {
  ByteBufferOutput wo=new ByteBufferOutput(cl.buf);

  wo.write_string("[ECM routes]");

  Hashtable seen=new Hashtable();

  print_established_ecm_routes(wo,seen);
  print_pending_ecm_routes(wo,seen);

  wo.write_int(-1);

  wo.flush();

  write_fully(console,wo.buf());
 }


 private static void process_go_msg(Client cl) throws Throwable {
  ByteBufferInput wi=new ByteBufferInput(cl.buf);
  int box=wi.get_byte();

  int res=1;

  if (box<clients.size()) {
   Client c=(Client)clients.elementAt(box);
   if (c!=console) {
    cur_box=c;
    res=0;
   }
  }

  ByteBufferOutput wo=new ByteBufferOutput(cl.buf);

  wo.write_int(res);

  wo.flush();

  write_fully(console,wo.buf());
 }

 private static void process_exit_msg(Client cl) throws Throwable {
  drop_client(console);
 }

 private static void process_console_data_msg(Client cl) throws Throwable {
  if (cur_box!=null) forward_msg(cl,cur_box);
   else discard_msg(cl);
 }

 private static void process_console_msg(Client cl) throws Throwable {
  switch(cl.msg_type) {
   case Proto.CMD_CLIENTS_COUNT:
    process_clients_count_msg(cl);
    break;
   case Proto.CMD_CLIENT_ID:
    process_client_id_msg(cl);
    break;
   case Proto.CMD_LIST:
    process_list_msg(cl);
    break;
   case Proto.CMD_GO:
    process_go_msg(cl);
    break;
   case Proto.CMD_EXIT:
    process_exit_msg(cl);
    break;
   case Proto.CMD_ECMROUTES:
    process_ecmroutes_msg(cl);
    break;
   case Proto.CMD_DATA:
    process_console_data_msg(cl);
    break;
  }
 }

 private static boolean collect_msg(Client cl) throws Throwable {
  buf.clear();

  int num_read=0;

  try {
   num_read=cl.ch.read(buf);
  } catch(IOException e) {
   return false;
  }

  if (num_read==-1) return false;

  buf.flip();

  while(buf.hasRemaining()) {  
   switch(cl.read_state) {
    case RS_READ_HEADER:
     byte b=buf.get();
     cl.hdr[cl.hdr_size++]=b;

     if (cl.hdr_size==HEADER_SIZE) {
      ByteArrayInputStream bais=new ByteArrayInputStream(cl.hdr);
      WrappedInput wi=new StreamInput(bais);

      int size=wi.read_short();
      int type=wi.get_byte();

      if (allowed_msg_type(type)&&box_msg_size(size)) {
       cl.read_state=RS_READ_DATA;
       cl.msg_type=type;
       cl.msg_size=size;
       cl.read_size=size;

       if (size==0) cl.has_msg=true;
      } else drop_client(cl);
     }
     break;
 
    case RS_READ_DATA:
     int cnt=buf.remaining();

     if (cl.read_size<=cnt) {
      int old_limit=buf.limit();
      int new_limit=buf.position()+cl.read_size;

      buf.limit(new_limit);
      cl.buf.put(buf);
      buf.limit(old_limit);

      cl.buf.flip();
 
      cl.has_msg=true;
     } else {
      cl.buf.put(buf);
      cl.read_size-=cnt;
     }
     break;
   }

   if (cl.has_msg) {
    switch(cl.type) {
     case Proto.CLT_BOX:
      process_box_msg(cl);
      break;

     case Proto.CLT_CONSOLE:
      process_console_msg(cl);
      break;
    }

    reset_msg(cl);
   }

  }

  return true;
 }

 private static void main_loop() {
  while(running) {
   try {
    selector.select();

    Set keys=selector.selectedKeys();
    Iterator i=keys.iterator();

    while(i.hasNext()) {
     SelectionKey key=(SelectionKey)i.next();

     i.remove();

     if (key.isAcceptable()) {
      SocketChannel cl_ch=srv_ch.accept();
      
      Socket sck=cl_ch.socket();

      InetAddress addr=sck.getInetAddress();

      byte tab[]=addr.getAddress();
      int ip=((tab[0]&0xff)<<24)|((tab[1]&0xff)<<16)|((tab[2]&0xff)<<8)|(tab[3]&0xff);

      InputStream is=sck.getInputStream();

      WrappedInput wi=new StreamInput(is);

      wi.get_byte();
      wi.get_byte();

      int type=read_client_type(wi);
      String desc=read_client_desc(wi);

      if ((type!=-1)&&(desc!=null)) {
       Client cl=new Client(type,ip,desc,cl_ch);

       add_client(cl);

       cl_ch.configureBlocking(false);
       cl_ch.register(selector,SelectionKey.OP_READ);

       System.out.println("new client: "+type+" "+desc);
      } else drop_channel(cl_ch);

      continue;
     }

     if (key.isReadable()) {
      SocketChannel cl_ch=(SocketChannel)key.channel();

      Client cl=find_client(cl_ch);

      if (cl!=null) {
       if (!collect_msg(cl)) {
        drop_client(cl);
        continue;
       }
      }

      buf.clear();
     }
    }
   } catch(Throwable t) {
    t.printStackTrace();
   }
  }
 }

 public static void main(String args[]) {
  int port=ConfigDefs.BOXSERVER_PORT;

  if (start_server(port)) {
   System.out.println("started BoxServer at "+port);
   main_loop();
  }
 }
}
