def testGraph():
    from neo4j.v1 import GraphDatabase, basic_auth
    driver = GraphDatabase.driver(graph_address, auth=basic_auth(graph_username, graph_password))
    session = driver.session()

    session.run("CREATE (a:Person {name: {name}, title: {title}})",
              {"name": "Arthur", "title": "King"})

    result = session.run("MATCH (a:Person) WHERE a.name = {name} "
                       "RETURN a.name AS name, a.title AS title",
                       {"name": "Arthur"})
    for record in result:
      print("%s %s" % (record["title"], record["name"]))

    session.close()

def runSSH(hostname, command):
    client = paramiko.Transport((hostname, ssh_port))
    client.connect(username=ssh_username, password=ssh_password)

    stdout_data = []
    stderr_data = []
    session = client.open_channel(kind='session')
    session.exec_command(command)
    while True:
        if session.recv_ready():
            stdout_data.append(session.recv(buffer_nbytes))
        if session.recv_stderr_ready():
            stderr_data.append(session.recv_stderr(buffer_nbytes))
        if session.exit_status_ready():
            break

    if session.recv_exit_status() == 0:
        return ''.join([str(x) for x in stdout_data])
    else:
        return 'exit status: '.join(session.recv_exit_status()).join(''.join([str(x) for x in stderr_data]))

    session.close()
    client.close()

def newSSH(hostname, command):
    client = paramiko.Transport((hostname, ssh_port))
    client.connect(username=ssh_username, password=ssh_password)

    stdout_data = []
    session = client.open_channel(kind='session')
    session.exec_command(command)
    while True:
        buf = session.recv(1024)
        if not buf:
            break
        stdout_data.append(buf)

    session.close()
    client.close()

    return ''.join([str(x) for x in stdout_data])

def testReadDatabase(input):
    import xml.etree.ElementTree as ET
    root = ET.fromstring(input)
    #for child in root[0][1]:
    xmlns = root[0].tag[root[0].tag.index('{'):root[0].tag.index('}')+1]
    print xmlns

    for child in root[0][1].findall(xmlns+'isis-database-entry'):
        #print child.tag, child.attrib
        hostname = child.find(xmlns+'isis-tlv') \
            .find(xmlns+'hostname-tlv') \
            .find(xmlns+'hostname').text
        print hostname
        ipaddress = child.find(xmlns+'isis-tlv') \
            .find(xmlns+'ipaddress-tlv') \
            .find(xmlns+'address').text
        print ipaddress

        #for child2 in tlv:
        #    print child2.tag#, child2.attrib
    #print db.rpc-reply.isis-database.isis-database-entry.isis-tlv.hostname-tlv
    #print child.tag, child.attrib

def main():
    #testGraph()
    print newSSH(ssh_hostname, 'show isis database extensive | display xml')
    #test = runSSH(ssh_hostname, 'show isis database extensive | display xml')
    #testReadDatabase(test)

    #print runSSH(ssh_hostname, 'show isis database | display xml')
    #print runSSH(ssh_hostname, 'show isis database GSHT-LAB-MC01.00-00 extensive | display xml')
    #print runSSH(ssh_hostname, 'show isis adjacency')
    #print runSSH(ssh_hostname, 'show isis hostname')
    #print runSSH(ssh_hostname, 'show isis interface')
    #print runSSH(ssh_hostname, 'show interfaces ae1')
    #print runSSH(ssh_hostname, 'show configuration interfaces ae1.0')
    #print runSSH(ssh_hostname, 'show configuration interfaces ae2.0')
    #print runSSH(ssh_hostname, 'show configuration interfaces lo0.0')
    #print runSSH(ssh_hostname, 'show isis database extensive | display xml')

    #nodes = getNodeList()
    #createNode(nodes)
    #createLink()

def getNodeList():
    input = runSSH(ssh_hostname, 'show isis database | display xml')
    root = ET.fromstring(input)

    xmlns = root[0].tag[root[0].tag.index('{'):root[0].tag.index('}')+1]
    #print xmlns

    result = []

    for child in root[0][1].findall(xmlns+'isis-database-entry'):
        #print child.tag, child.attrib
        lspid = child.find(xmlns+'lsp-id').text
        if not lspid in result:
            result.append(lspid)

    return result

def createNode(nodes):
    from neo4j.v1 import GraphDatabase, basic_auth
    driver = GraphDatabase.driver(graph_address, auth=basic_auth(graph_username, graph_password))
    session = driver.session()

    session.run("MATCH (n:Node) OPTIONAL MATCH (n)-[r]-() DELETE r,n")

    for node in nodes:
        input = runSSH(ssh_hostname, 'show isis database %s extensive | display xml' % (node))
        root = ET.fromstring(input)

        xmlns = root[0].tag[root[0].tag.index('{'):root[0].tag.index('}')+1]
        #print xmlns

        for child in root[0][1].findall(xmlns+'isis-database-entry'):
            #print child.tag, child.attrib
            hostname = child.find(xmlns+'isis-tlv') \
                .find(xmlns+'hostname-tlv') \
                .find(xmlns+'hostname').text
            ipaddress = child.find(xmlns+'isis-tlv') \
                .find(xmlns+'ipaddress-tlv') \
                .find(xmlns+'address').text
            print hostname
            print ipaddress
            print '----------'
            sys.stdout.flush()

            session.run("CREATE (a:Node {title: {title}, hostname: {hostname}, ipaddress: {ipaddress}})",
                       {"title": hostname, "hostname": hostname, "ipaddress": ipaddress})

    session.close()

def createLink():
    from neo4j.v1 import GraphDatabase, basic_auth
    driver = GraphDatabase.driver(graph_address, auth=basic_auth(graph_username, graph_password))
    session = driver.session()

    nodes = session.run("MATCH (n:Node) RETURN n.hostname AS hostname, n.ipaddress AS ipaddress")

    for node in nodes:
        hostname = node["hostname"]
        ipaddress = node["ipaddress"]

        input = runSSH(ipaddress, 'show isis adjacency | display xml')
        root = ET.fromstring(input)

        xmlns = root[0].tag[root[0].tag.index('{'):root[0].tag.index('}')+1]
        #print xmlns

        for child in root[0].findall(xmlns+'isis-adjacency'):
            #print child.tag, child.attrib
            interface_name = child.find(xmlns+'interface-name').text
            system_name = child.find(xmlns+'system-name').text
            adjacency_state = child.find(xmlns+'adjacency-state').text
            print hostname
            print system_name
            print interface_name
            print adjacency_state
            print '----------'
            sys.stdout.flush()

            command = "MATCH (n1:Node {hostname:'%s'}), \
                               (n2:Node {hostname:'%s'}) \
                        CREATE (n1)-[:LINKED {name:'%s', state:'%s'}]->(n2)" % \
                        (hostname, system_name, interface_name, adjacency_state)
            session.run(command)
            #session.run("MATCH (n1:Node {hostname:'{hostname1}'}), \
            #                   (n2:Node {hostname:'{hostname2}'}) \
            #            CREATE (n1)-[:LINKED {name:'{name}', state:'{state}'}]->(n2)",
            #            {"hostname1": hostname, "hostname2": system_name, "name": interface_name, "state": adjacency_state})

    session.close()

if __name__ == "__main__" :
    main()