from metasploit.msfrpc import MsfRpcClient
import Exploiter.exploits as exploits
import Exploiter.postexploits as pexploits

def callExploit(client:MsfRpcClient, exploitName:str, targetIP:str, localIP:str):
    exploitMap = exploits.exploitMap
    return exploitMap[exploitName](client, targetIP, localIP)

def callPostExploit(client:MsfRpcClient, postExploitAction:str, sessionNum:int):
    postExploitMap = pexploits.postexploitMap
    return postExploitMap[postExploitAction](client, sessionNum)

def callPostExploitByIP(client:MsfRpcClient, postExploitAction:str, targetIP:str):
    sessions = client.sessions.list
    for session in sessions.keys():
        if sessions[session]["target_host"] == targetIP:
            return callPostExploit(client, postExploitAction, int(session))

def callPostExploitByUUID(client:MsfRpcClient, postExploitAction:str, sessionUUID:str):
    sessions = client.sessions.list
    for session in sessions.keys():
        if sessions[session]["uuid"] == sessionUUID:
            return callPostExploit(client, postExploitAction, int(session))

def getSessionbyIP(client:MsfRpcClient, targetIP:str):
    sessions = client.sessions.list
    for session in sessions.keys():
        if sessions[session]["target_host"] == targetIP:
            return session

def getSessionbyUUID(client:MsfRpcClient, sessionUUID:str):
    sessions = client.sessions.list
    for session in sessions.keys():
        if sessions[session]["uuid"] == sessionUUID:
            return session

def getSessionbyExploitUUID(client:MsfRpcClient, exploitUUID:str):
    sessions = client.sessions.list
    for session in sessions.keys():
        if sessions[session]["exploit_uuid"] == exploitUUID:
            return session

def getExploitNames():
    return exploits.exploitMap.keys()
