module Tri

# File lib/tri.rb, line 34
def self.get_client
        if is_running?
                client = @server.accept
                
                # Got a connection from a client
                if key = validWebSocketRequest?(client)
                        # Send the respone back for handshaking phase
                        client.print  caculateRequestReply(key)    
                else
                        # Close the connection due to invalid WebSocket requests
                        client.close
                        client = nil     
                end
        end
        
        # Return the client
        client
end

# File lib/tri.rb, line 54
def self.get_client_with_keep_alive(second)
        if is_running?
                client = @server.accept
                
                # Got a connection from a client
                if key = validWebSocketRequest?(client)
                        # Send the respone back for handshaking phase
                        client.print  caculateRequestReply(key)
                         # Sending PING, for every 28 sec
                        Thread.new {
                                while client
                                        ping_to(client)
                                        sleep(second)           
                                end
                        }
                else
                        # Close the connection due to invalid WebSocket requests
                        client.close
                        client = nil     
                end
        end
        
        # Return the client
        client
end

# File lib/tri.rb, line 29
def self.is_running?
        !@server.nil?
end

# File lib/tri.rb, line 100
def self.ping_to(client)
    pingData    = "CHAT4NINJA"
    firstByte   = "10001001".to_i(2).chr
    secondByte  = pingData.size.chr
    client.write(firstByte + secondByte + pingData)
end

# File lib/tri.rb, line 81
def self.receive_from(client)
        data   = getData(client)
        case data
            when -2       
                [-2,"[ERROR] Huge payload - should close the connection."]
            when -1       
                [-2,"[ERROR] No data due to connection timeout."]
            when 8       
                [8, "[CLOSE] Close frame sent from the client."]
            when 9       
                [8, "[OPCODE] PING frame sent from the client."]   
            when 10       
                [9, "[OPCODE] PONG frame sent from the client."]        
            else
                [1, data] 
        end
end

# File lib/tri.rb, line 108
def self.send_to(client,data)
        firstByte   = "10000001".to_i(2).chr
        if data.size <=127
            secondByte  = data.size.chr
            client.write(firstByte + secondByte + data)
        else
            binaryString= data.size.to_s(2)
            extraZero   = ""
            (16 - binaryString.size).times do
                    extraZero += "0"
            end
            extraZero  = extraZero + binaryString 
            secondByte = "01111110".to_i(2).chr
            thirdByte  = extraZero[0..7].to_i(2).chr
            forthByte  = extraZero[8..15].to_i(2).chr
       
            begin
                client.write(firstByte + secondByte + thirdByte + forthByte + data)
                # Catch err raised by the command above
                rescue SystemCallError => err
                       # Return error code
                       "[ERROR] Trying to send {" + firstByte + secondByte + thirdByte + forthByte + data + "} " + err.to_s 
                else
                       # Return success message
                      
            end
        end
end

# File lib/tri.rb, line 6
def self.start(input1="localhost", input2 = 888)
        # Check input parameters
        if !input1.is_a? String
                hostname = "localhost"
                port     = input1
        else
                hostname = input1
                port     = input2
        end
        # Trying to init a new TCPserver
        begin
                @server = TCPServer.new(hostname,port)
                # Catch err raised by the command above
                rescue SystemCallError => err
                       # Return error code
                       "[ERROR] Trying to start the server @"  + hostname + ":" + port.to_s + ".\r\n[ERROR] "  + err.to_s 
                else
                       # Return success message
                       "[SERVER] The server is now listening @" + hostname + ":" + port.to_s + "."
        end
end

# File lib/tri.rb, line 158
def self.caculateRequestReply(key)
      "HTTP/1.1 101 Switching Protocols\r\n" +
      "Upgrade: websocket\r\n" +
      "Connection: Upgrade\r\n" +
      "Sec-WebSocket-Accept: #{sha_1(key)}\r\n\r\n" 
end

# File lib/tri.rb, line 171
def self.decodeData(input, mask)
    data = ""
    for i in 0..7 do
        if input[i]==mask[i]
            data += "0"
        else
            data += "1"
        end
    end
    data.to_i(2).chr
end

# File lib/tri.rb, line 184
def self.getData(client)
        frame           = Hash.new
        # First 8 bits after connected
        receiveByte = client.recv(1).unpack("B*").join
        if receiveByte != ""
            
            frame["fin"]    = receiveByte[0].to_i(2)
            frame["rsv"]    = receiveByte[1..3].to_i(2)
            frame["opcode"] = receiveByte[4..7].to_i(2)
            
            # Second 8 bits
            receiveByte     = client.recv(1).unpack("B*").join
            frame["mask"]   = receiveByte[0].to_i(2)
            frame["payload"]= receiveByte[1..7].to_i(2)
            
            
            # Load more frame["payload"]?
            if frame["payload"] >= 126 
                loadMore    = 0
                newPayload  = ""
                if  frame["payload"] == 126
                         loadMore = 2
                end
                if  frame["payload"] == 127
                         loadMore = 8
                end
                loadMore.times do
                    newPayload += client.recv(1).unpack("B*").join 
                end
                frame["payload"] = newPayload.to_i(2)
            end
            
            if frame["payload"] <= 500
                    # Load the MASKEY or not
                    if frame["mask"] == 1
                        frame["mask-key"] = Array.new
                        for i in 0..3
                                frame["mask-key"].push client.recv(1).unpack("B*").join    
                        end
                    end
                    
                    # Load data based-on the payload
                    frame["data"] = ""
                    for i in 0..(frame["payload"]-1)
                            frame["data"] += decodeData(client.recv(1).unpack("B*").join,frame["mask-key"][i % 4])
                    end
                    
                    if frame["opcode"] != 1
                            frame["opcode"]
                    else    
                            frame["data"]
                    end
            else
                    -2 
            end
        else
            -1  # no data
        end
end

# File lib/tri.rb, line 166
def self.sha_1(input)
    Digest::SHA1.base64digest(input+"258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
end

# File lib/tri.rb, line 140
def self.validWebSocketRequest?(client)
    handshakeKeys   = Hash.new
    if (client.gets.include?("GET / HTTP/1.1"))
          while  (requests = client.gets) && (requests != "\r\n")
            if requests.include?(": ")
                  requests = requests.split(": ")
                  if requests.size == 2
                        handshakeKeys[requests[0]] = requests[1].chomp
                  end      
            end      
          end
          handshakeKeys["sec-websocket-key"]            # return the key
    else
          ""                                            # return nil
    end
end