ging/social_stream

#!/usr/bin/env ruby

####################################
#          EMANAGEMENT
#Ejabberd Management script 
#New features for management and maintenance ejabberd
#@author Aldo Gordillo < agordillos@gmail.com >
#@version 2.1 - 28-2-2012
####################################

require 'logger'

path = "/var/log/ejabberd/scripts.log"
file = File.open(path, File::WRONLY | File::APPEND | File::CREAT)
file.sync = true
$logger = Logger.new(file)
$logger.level = Logger::DEBUG

def getOption(option)
  File.open('/etc/ejabberd/ssconfig.cfg', 'r') do |f1|  
    while line = f1.gets  
      line = line.gsub(/\n/,'')
      if line.match(/^#/)
        #Comments
      elsif line.match(/^#{option}/)
        return line.gsub(/#{option}/,'')
      end  
    end  
  end
  return "Undefined"
end


#Configuration variables
$verbose = (getOption("verbose=")=="true")
$ejabberd_user = getOption("ejabberd_server_user=")
$checkEjabberdctlQuotedString = false


PARAMS_FOR_COMMANDS = {
    'addBuddyToRoster' => 5,
    'removeBuddyFromRoster' => 2,
    'setBidireccionalBuddys' => 6,
    'unsetBidireccionalBuddys' => 4,
    'getRoster' => 1,
    'removeRoster' => 1,
    'removeAllRostersByDomain' => 1,
    'removeAllRosters' => 0,
    'getAllUserJidsWithRosterByDomain' => 1,
    'printAllRostersByDomain' => 1,
    'printAllRosters' => 0,
    'printAllBidirecctionalBuddysByDomain' => 1,
    'checkUserJid' => 1,
    'checkBidirecctionalBuddys' => 2,
    'sendPresence' => 2,
    'setPresence' => 1,
    'unsetPresence' => 1,
    'sendMessageToUser' => 3,
    'getUserResource' => 1,
    'isEjabberdNodeStarted' => 0,
    'broadcast' => 3,
    'broadcastToConnectedUsers' => 3,
        'getConnectedJidsByDomain' => 1,
    'getConnectedJids' => 0,
    'kickUserJid' => 1,
    'createPersistentRoom' => 2,
    'createRoom' => 2,
    'destroyRoom' => 2,
    'destroyAllRoomsByDomain' => 1,
    'destroyAllRooms' => 0,
    'getAllJidsOfRoom' => 2,
    'getAffiliationsOfRoom' => 2,
    'setJidAffiliationOfRoom' => 4,
    'printAllRoomsByDomain' => 1,
    'printAllRooms' => 0,
    'help' => 0,
}

SYNTAX_FOR_COMMANDS = {
    'addBuddyToRoster' => 'addBuddyToRoster userJid buddyJid buddyNick buddyGroup subscription_type',
    'removeBuddyFromRoster' => 'removeBuddyFromRoster userJid buddyJid',
    'setBidireccionalBuddys' => 'setBidireccionalBuddys userAJid userBJid userANick userBNick groupForA groupForB',
    'unsetBidireccionalBuddys' => 'unsetBidireccionalBuddys userJid oldFriendJid oldFriendNick groupForOldFriend',
    'getRoster' => 'getRoster userJid',
    'removeRoster' => 'removeRoster userJid',
    'removeAllRostersByDomain' => 'removeAllRostersByDomain domain',
    'removeAllRosters' => 'removeAllRosters',
    'getAllUserJidsWithRosterByDomain' => 'getAllUserJidsWithRosterByDomain domain',
    'printAllRostersByDomain' => 'printAllRostersByDomain domain',
    'printAllRosters' => 'printAllRosters',
    'printAllBidirecctionalBuddysByDomain' => 'printAllBidirecctionalBuddysByDomain domain',
    'checkUserJid' => 'checkUserJid userJid',
    'checkBidirecctionalBuddys' => 'checkBidirecctionalBuddys userAJid userBJid',
    'sendPresence' => 'sendPresence userJid show',
    'setPresence' => 'setPresence userJid',
    'unsetPresence' => 'unsetPresence userJid',
    'sendMessageToUser' => 'sendMessageToUser fromJid toJid msg',
    'getUserResource' => 'getUserResource userJid',
    'isEjabberdNodeStarted' => 'isEjabberdNodeStarted',
    'broadcast' => 'broadcast admin userJids msg ',
    'broadcastToConnectedUsers' => 'broadcastToConnectedUsers admin userJids msg (userJids value: "all" or jids array [jid1,jid2,...,jidN])',
        'getConnectedJidsByDomain' => 'getConnectedJidsByDomain domain',
    'getConnectedJids' => 'getConnectedJids',
    'kickUserJid' => 'kickUserJid(userJid)',
    'createPersistentRoom' => 'createPersistentRoom roomName domain',
    'createRoom' => 'createRoom roomName domain',
    'destroyRoom' => 'destroyRoom roomName domain',
    'destroyAllRoomsByDomain' => 'destroyAllRoomsByDomain domain',
    'destroyAllRooms' => 'destroyAllRooms',
    'getAllJidsOfRoom' => 'getAllJidsOfRoom roomName domain',
    'getAffiliationsOfRoom' => 'getAffiliationsOfRoom roomName domain',
    'setJidAffiliationOfRoom' => 'setJidAffiliationOfRoom roomName domain jid affiliation (affiliation value: owner/admin/member/outcast/none)',
    'printAllRoomsByDomain' => 'printAllRoomsByDomain domain',
    'printAllRooms' => 'printAllRooms',
    'help' => 'help',
}


#########################
#Debug methods
#########################
def ejabberdLog(text)
    $logger.info "Ejabberd Management Script: " + text
    if $verbose
        #puts "Writing to ejabberdLog: " + "Ejabberd Management Script: " + text
    end
end

def log(msg)
    logWithTitle(msg,nil)
end

def logWithTitle(msg,title)
    puts "------------------------"
    if title
        puts ("<<<" + title + ">>>")
    end
    puts msg
    puts "------------------------"
end


#########################
#Methods to manage rosters from Social Stream Rails App
#########################
def setBidireccionalBuddys(userAJid,userBJid,userANick,userBNick,groupForA,groupForB)
    addBuddyToRoster(userAJid,userBJid,userBNick,groupForB,"both")
    addBuddyToRoster(userBJid,userAJid,userANick,groupForA,"both")
    return "Done"
end

def unsetBidireccionalBuddys(userJid,oldFriendJid,oldFriendNick,groupForOldFriend)
    if checkBidirecctionalBuddys(userJid,oldFriendJid)
        removeBuddyFromRoster(userJid,oldFriendJid)
        removeBuddyFromRoster(oldFriendJid,userJid)
        addBuddyToRoster(userJid,oldFriendJid,oldFriendNick,groupForOldFriend,"to")
        return "Done"
    else
        return userJid + " and " + oldFriendJid + " aren't bidireccional buddys"
    end
end

def addBuddyToRoster(userJid,buddyJid,buddyNick,buddyGroup,subscription_type)
    user = getUsernameFromJid(userJid)
    buddy = getUsernameFromJid(buddyJid)
    userDomain = getDomainFromJid(userJid)
    buddyDomain = getDomainFromJid(buddyJid)
    executeCommand("ejabberdctl add-rosteritem " + user + " " + userDomain + " " + buddy + " " + buddyDomain + " " + buddyNick + " " + buddyGroup + " " +                 subscription_type)
    return "Done"
end

def removeBuddyFromRoster(userJid,buddyJid)
    user = getUsernameFromJid(userJid)
    buddy = getUsernameFromJid(buddyJid)
    userDomain = getDomainFromJid(userJid)
    buddyDomain = getDomainFromJid(buddyJid)
    if checkUserJidInRoster(buddyJid,getRoster(userJid))
        executeCommand("ejabberdctl delete_rosteritem " + user + " " + userDomain + " " + buddy + " " + buddyDomain)
        return "Done"
    else
        return "User " + buddyJid + " not found in " + userJid + " roster."
    end
end


#########################
#Roster Utilities
#########################
def getRoster(userJid)
    if checkUserJid(userJid)
        executeCommand("ejabberdctl get_roster " + getUsernameFromJid(userJid)  + " " + getDomainFromJid(userJid))
    else
        return "Roster not found for user " + userJid
    end
end

def getBuddyJidsFromRoster(roster)
    buddyJids = []
        lines = roster.split("\n")
    lines.each do |line|
        buddyJids << line.split(" ")[0]
    end
    buddyJids
end

def removeRoster(userJid)
    if checkUserJid(userJid)
        executeCommand("ejabberdctl process_rosteritems delete any any " + userJid + " any")
        return "Done"    
    else
        return "Roster not found for userJid " + userJid
    end
end

def removeAllRostersByDomain(domain)
    if(domain=="all")
        executeCommand("ejabberdctl process_rosteritems delete any any any any")
    else
        executeCommand("ejabberdctl process_rosteritems delete any any *@" + domain + " any")
    end
    return "Done";
end

def removeAllRosters()
    return removeAllRostersByDomain("all")
end

def getAllUserJidsWithRosterByDomain(domain)
    if(domain=="all")
        output = executeCommand("ejabberdctl process_rosteritems list any any any any")
    else
        output = executeCommand("ejabberdctl process_rosteritems list any any *@" + domain + " any")
    end

    userJids = []
    lines = output.split("\n");
    items = lines[0].split(" ")[2]

    #test if items string contains a correct number
    if items.to_i.to_s == items
        if items.to_i > 0
            lines.each do |line|
                if line.split(":")[0]=="Matches"
                    userJid = line.split(" ")[1]
                    #puts "Line:" + line
                    unless userJids.include?(userJid)
                        userJids << userJid
                    end    
                end
            end
        end
    end
    
    return userJids;
end

def getAllUserJidsWithRoster()
    return getAllUserJidsWithRosterByDomain("all")
end

def getAllRostersByDomain(domain)
    rosterList = { }

    userJids = getAllUserJidsWithRosterByDomain(domain)

    userJids.each do |userJid|
        roster = getRoster(userJid)
        rosterList.store(userJid,roster)
    end
    return rosterList
end

def getAllRosters()
    return getAllRostersByDomain("all");
end

def printAllRostersByDomain(domain)
    rosterList = getAllRostersByDomain(domain)
    rosterList.keys.each do |userJid|
        puts "\n"
        puts "-------------------------------------"
        puts userJid + " Roster"
        puts "-------------------------------------"
        puts rosterList[userJid]
        puts "-------------------------------------"
        puts "\n"
    end
    return "Done"
end

def printAllRosters()
    return printAllRostersByDomain("all")
end

def getAllBidirecctionalBuddysByDomain(domain)
    b_buddyJids = []
    userJids = getAllUserJidsWithRosterByDomain(domain)
    nonCheckedUsers = userJids
    userJids.each do |userJid|
        nonCheckedUsers.delete(userJid)
        nonCheckedUsers.each do |checkUserJid|
            if checkBidirecctionalBuddys(userJid,checkUserJid)
                b_buddyJids << [userJid,checkUserJid]
            end
        end
    end
    return b_buddyJids
end

def getAllBidirecctionalBuddys()
    return getAllBidirecctionalBuddysByDomain("all")
end

def printAllBidirecctionalBuddysByDomain(domain)
    puts "This may take a while..."
    b_buddys = getAllBidirecctionalBuddysByDomain(domain)
    b_buddys.each do |contact|
        puts "[" + contact[0] + "," + contact[1] + "]"
    end
    return "Done"
end

def printAllBidirecctionalBuddys()
    return printAllBidirecctionalBuddysByDomain("all")
end

#Check if the user have a roster in his domain
def checkUserJid(userJid)
    domain = getDomainFromJid(userJid)
    return getAllUserJidsWithRosterByDomain(domain).include?(userJid)
end

def checkUserJidInRoster(userJid,roster)
    return getBuddyJidsFromRoster(roster).include?(userJid)
end

def checkBidirecctionalBuddys(userAJid,userBJid)
    rosterA = getRoster(userAJid)
    rosterB = getRoster(userBJid)
    return (checkUserJidInRoster(userAJid,rosterB) and checkUserJidInRoster(userBJid,rosterA))
end


#########################
#Manage stanzas Utilities
#########################
def setPresence(userJid)
    sendPresenceStanzaWithType(userJid,userJid,"available")
end

def unsetPresence(userJid)
    sendPresenceStanzaWithType(userJid,userJid,"unavailable")
end

def sendPresence(username,show)
    sendPresenceWithShow(username,username,show)
end

def sendPresenceWithShow(fromJid,toJid,show)
    resource = getUserResource(fromJid);  
    from_name = getUsernameFromJid(fromJid)
    domain = getDomainFromJid(fromJid)
    pres_stanza = "\\<" + buildQuotedString("presence from=") + "\\\"" + buildQuotedString(fromJid) + "\\\"" + buildQuotedString(" to=") + "\\\"" +
            buildQuotedString(toJid) + "\\\"\\>\\<" + buildQuotedString("show") + "\\>" + buildQuotedString(show) + "\\<" +
            buildQuotedString("/show") + "\\>\\<" + buildQuotedString("/presence") + "\\>"
    executeCommand("ejabberdctl send_stanza_c2s " + from_name + " " + domain + " " + resource + " " + pres_stanza)
    return "Done"
end

def sendPresenceStanzaWithType(fromJid,toJid,presence_type)
    resource = getUserResource(fromJid);
    from_name = getUsernameFromJid(fromJid)
    domain = getDomainFromJid(fromJid)
    pres_stanza = "\\<" + buildQuotedString("presence type=") + "\\\"" + buildQuotedString(presence_type) + "\\\"" + buildQuotedString(" from=") + "\\\"" + 
        buildQuotedString(fromJid) + "\\\"" + buildQuotedString(" to=") + "\\\"" + buildQuotedString(toJid) + "\\\"\\>\\<" + buildQuotedString("/presence") + "\\>" 
    executeCommand("ejabberdctl send_stanza_c2s " + from_name + " " + domain + " " + resource + " " + pres_stanza)
    return "Done"
end

def sendMessageToUser(fromJid,toJid,msg)
    executeCommand("ejabberdctl send_message_chat " + fromJid + " " + toJid + " " + buildQuotedString(msg))
    return "Done"
end

def getUserResource(userJid)
    output = executeCommand("ejabberdctl connected-users")
        lines = output.split("\n");
    lines.each do |line|
        if line.split("/")[0] == userJid
            #puts "Find " + userJid
            resource = line.split("/")[1];
            return resource;    
        end
    end
    return userJid + " no have any active session"
end


#########################
#Utilities
#########################
def getUsernameFromJid(jid)
    return jid.split("@")[0];
end

def getDomainFromJid(jid)
    return jid.split("@")[1];
end

def getElementsFromStringArray(stringArray)
    stringArray=stringArray[1,stringArray.length-2]
    return stringArray.split(",")
end


#Determine how to scape characters for build quoted strings
def checkEjabberdctlQuotedString
    puts "checkForSimpleSlash: " + checkForSimpleSlash.to_s()
    puts "checkForDoubleSlash: " + checkForDoubleSlash.to_s()
end

def checkAndSetEjabberdctlQuotedString
    if checkForSimpleSlash
        $checkForSimpleSlash = true
    end
    if checkForDoubleSlash
        $checkForDoubleSlash = true
    end
    $checkEjabberdctlQuotedString = true    
end

def checkForDoubleSlash
    command = "ejabberdctl send_message_chat example@localhost example@localhost \\'Hello quoted string\\'"
    if execute_as_sudo
        command = "sudo " + command
    end
    
    output = %x[#{command}]
    firstLine = ""
    lines = output.split("\n")
    lines.each do |line|
        if line != ""
            firstLine = line
            break
        end
    end

    if firstLine==""
        return true
    elsif firstLine.split(":")[0]=="Error"
        return false
    else
        #Unknown error        
        return false
    end
end

def checkForSimpleSlash
    command = "ejabberdctl send_message_chat example@localhost example@localhost \'Hello quoted string\'"
    if execute_as_sudo
        command = "sudo " + command
    end
    #puts "Executing " + command
    output = %x[#{command}]
    firstLine = ""
    lines = output.split("\n")
    lines.each do |line|
        if line != ""
            firstLine = line
            break
        end
    end

    if firstLine==""
        return true
    elsif firstLine.split(":")[0]=="Error"
        return false
    else
        #Unknown error        
        return false
    end
end

def buildQuotedString(msg)
    if !$checkEjabberdctlQuotedString
        checkAndSetEjabberdctlQuotedString
    end
    
    if $checkForSimpleSlash
        return "\'" + msg + "\'"
    end
    if $checkForDoubleSlash
        return "\\'" + msg + "\\'"
    end
    return msg
end


#########################
#Connection Info
#########################
def isEjabberdNodeStarted
    output = executeCommand("ejabberdctl status")
    if firstLine = output.split("\n")[0]
      return ((firstLine.split(":")[1]).strip()=="started")
    end
    return false
end

def getConnectedJidsByDomain(domain)
        jids = []
    if(domain=="all")
        output = executeCommand("ejabberdctl connected-users")
    else
        output = executeCommand("ejabberdctl connected-users | grep @" + domain + "/")
    end
    sessions = output.split("\n")
    sessions.each do |session|  
      jids << session.split("/")[0]
    end
    return jids
end

def getConnectedJids()
    return getConnectedJidsByDomain("all")
end


#########################
#Advanced features
#########################
def broadcast(admin,userJids,msg)
    getElementsFromStringArray(userJids).each do |userJid|
        sendMessageToUser(admin,userJid,msg)
    end
    return "Done"
end

def broadcastToConnectedUsers(admin,userJids,msg)
    connectedJids = getConnectedJids()
    if(userJids=="all")
        connectedJids.each do |userJid|
            sendMessageToUser(admin,userJid,msg)
        end
    else
        getElementsFromStringArray(userJids).each do |userJid|
            if (connectedJids.include?(userJid))
                sendMessageToUser(admin,userJid,msg)
            end
        end
    end
    return "Done"
end

def kickUserJid(userJid)
    user = getUsernameFromJid(userJid)
    userDomain = getDomainFromJid(userJid)
    resource = getUserResource(userJid);  
    reason = "0"
    executeCommand("ejabberdctl kick_session " + user + " " + userDomain + " " + resource + " " + reason)
    return "Done"
end

def banUserJid(userJid)
    kickUserJid(userJid)
    #Notify Web server to ban user permanently
end



#########################
#MUC Methods (Multi-user-chat)
#########################

$muc_host = "conference"

def printAllRoomsByDomain(domain)
    if(domain=="all")
        return printAllRooms
    end
    output = executeCommand("ejabberdctl muc_online_rooms " + domain)
    return output;
end

def printAllRooms
    output = executeCommand("ejabberdctl muc_online_rooms global")
    return output;
end

def createPersistentRoom(roomName,domain)
    createRoom(roomName,domain)
    setRoomPersistence(roomName,domain,true)
    return "Done"
end

def createRoom(roomName,domain)
    executeCommand("ejabberdctl create_room " + roomName + " " + $muc_host + "." + domain + " " + domain)
    return "Done"
end

def setRoomPersistence(roomName,domain,persistence)
    if (persistence==true)
        executeCommand("ejabberdctl change_room_option " + roomName + " " + $muc_host + "." + domain + " persistent true")
    else
        executeCommand("ejabberdctl change_room_option " + roomName + " " + $muc_host + "." + domain + " persistent false")
    end
end

def destroyAllRoomsByDomain(domain)
    rooms = getAllRoomsByDomain(domain)
    rooms.each do |room|
        destroyRoom(getNameFromRoom(room),getDomainFromRoom(room))
    end
    return "Done"
end

def destroyAllRooms()
    return destroyAllRoomsByDomain("all")
end

def destroyRoom(roomName,domain)
    executeCommand("ejabberdctl destroy_room " + roomName + " " + $muc_host + "." + domain + " " + domain)
    return "Done"
end

def getAllRoomsByDomain(domain)
    if(domain=="all")
        output = executeCommand("ejabberdctl muc_online_rooms global")
    else
        output = executeCommand("ejabberdctl muc_online_rooms " + domain)
    end
    rooms = output.split("\n")
    return rooms
end

def getAllJidsOfRoomObject(room)
    return getAllJidsOfRoom(getNameFromRoom(room),getDomainFromRoom(room))
end

def getAllJidsOfRoom(roomName,domain)
    jids = []
    output = executeCommand("ejabberdctl get_room_occupants " + roomName + " " +  $muc_host + "." + domain)
    lines = output.split("\n")
    lines.each do |line|
        jids << line.split("/")[0]
    end
    return jids
end

def getAffiliationsOfRoomObject(room)
    return getAffiliationsOfRoom(getNameFromRoom(room),getDomainFromRoom(room))
end

def getAffiliationsOfRoom(roomName,domain)
    affiliateds = []
    output = executeCommand("ejabberdctl get_room_affiliations " + roomName + " " +  $muc_host + "." + domain)
    lines = output.split("\n")
    lines.each do |line|
        split=line.split(" ");
        jid = split[0]+"@"+split[1];
        affiliation=split[2]
        affiliateds << [jid,affiliation]
    end
    return affiliateds
end

#affiliation = Owner/Admin/Member/Outcast/None
def setJidAffiliationOfRoom(roomName,domain,jid,affiliation)
    output = executeCommand("ejabberdctl set_room_affiliation " + roomName + " " +  $muc_host + "." + domain + " " + jid + " " + affiliation)
    return "Done"
end

def getDomainFromRoom(room)
    return room.split(".")[1];
end

def getNameFromRoom(room)
    return room.split("@")[0];
end


#########################
#Execution commands methods
#########################
def executeCommand(command)
    #Building...
    command = buildCommand(command)
    
    if $verbose
        #Logging...
        puts "Executing: " + command
        ejabberdLog("Executing (#{command})")
    end

    #Executing...
    output = %x[#{command}]
    return output
end

def buildCommand(command)
    if execute_as_sudo
        command = "sudo -u " + $ejabberd_user + " " + command
    end
    return command
end

def execute_as_sudo
    sudo_users = getOption("users_require_sudo=")

    if sudo_users=="all"
     return true
    end

    sudo_users_array = sudo_users.split(",")
        current_user = %x["whoami"].split("\n")[0]
    if sudo_users_array.include?(current_user)
        return true
    end
end


#########################
#Support
#########################
def help
    log("Command list")
    SYNTAX_FOR_COMMANDS.values.each do |command|
        puts command
    end
    puts ""
end



#########################
#Main thread
#########################
#log("Init Ejabberd Maintenance script")

begin
    if ARGV[0] and PARAMS_FOR_COMMANDS.keys.include?(ARGV[0])
        if (ARGV.length == (PARAMS_FOR_COMMANDS[ARGV[0]]+1))

            ejabberdLog("Executing (#{ARGV})")

            length = ARGV.length;
            case length
            when 1
              puts send(ARGV[0])
            when 2
              puts send(ARGV[0],ARGV[1])
            when 3
              puts send(ARGV[0],ARGV[1],ARGV[2])
            when 4
              puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3])
            when 5
              puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3],ARGV[4])
            when 6
              puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3],ARGV[4],ARGV[5])
            when 7
              puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3],ARGV[4],ARGV[5],ARGV[6])
            else
              puts send(ARGV[0],ARGV)
            end
            puts ""
        else 
            puts "Error: Params required for command " + (ARGV[0]).to_s() + ": " + (PARAMS_FOR_COMMANDS[ARGV[0]]).to_s()
            puts "Syntax for command " + (ARGV[0]).to_s() + "\n" + (SYNTAX_FOR_COMMANDS[ARGV[0]]).to_s()
            puts "Use 'help' to get information about command syntax"
        end
    else
        if (ARGV.length==0)
            help
        else
            puts "Error: Command not recognized"
            puts "Use 'help' to get information about command syntax"
        end
    end
rescue
    puts "Syntax error"
    puts "Use 'help' to get information about command syntax"
end