class Console(object):
def echo(self, text, file=None, nl=True, err=False, color=None, **styles):
return click.secho(text,
file=file,
nl=nl,
err=err,
color=color,
**styles)
def prompt(self, *args, **kwargs):
return prompt(*args, **kwargs)
multi_line = False
watcher = None
tx_colour = "yellow"
err_colour = "reset"
meta_colour = "cyan"
prompt_colour = "cyan"
def __init__(self, uri=None, **settings):
self.output_file = settings.pop("file", None)
verbose = settings.pop("verbose", False)
connection_data = get_connection_data(uri, **settings)
try:
self.graph = Graph(uri, **settings)
except ServiceUnavailable as error:
raise ConsoleError("Could not connect to {} -- {}".format(
connection_data["uri"], error))
try:
makedirs(HISTORY_FILE_DIR)
except OSError:
pass
self.history = FileHistory(path_join(HISTORY_FILE_DIR, HISTORY_FILE))
self.prompt_args = {
"history":
self.history,
"lexer":
PygmentsLexer(CypherLexer),
"style":
style_from_pygments(
VimStyle, {
Token.Prompt:
"#ansi{}".format(self.prompt_colour.replace(
"cyan", "teal")),
Token.TxCounter:
"#ansi{} bold".format(
self.tx_colour.replace("cyan", "teal")),
})
}
self.lexer = CypherLexer()
self.result_writer = Table.write
if verbose:
from neo4j.util import watch
self.watcher = watch("neo4j.%s" % connection_data["scheme"])
self.commands = {
"//": self.set_multi_line,
"/e": self.edit,
"/?": self.help,
"/h": self.help,
"/help": self.help,
"/x": self.exit,
"/exit": self.exit,
"/play": self.play,
"/csv": self.set_csv_result_writer,
"/table": self.set_tabular_result_writer,
"/tsv": self.set_tsv_result_writer,
"/config": self.config,
"/kernel": self.kernel,
}
self.tx = None
self.tx_counter = 0
def loop(self):
self.echo(TITLE, err=True)
self.echo("Connected to {}".format(self.graph.database.uri).rstrip(),
err=True)
self.echo(u"", err=True)
self.echo(dedent(QUICK_HELP), err=True)
while True:
try:
source = self.read()
except KeyboardInterrupt:
continue
except EOFError:
return 0
try:
self.run(source)
except ServiceUnavailable:
return 1
def run_all(self, sources):
gap = False
for s in sources:
if gap:
self.echo("")
self.run(s)
if not is_command(s):
gap = True
return 0
def run(self, source):
source = source.strip()
if not source:
return
try:
if is_command(source):
self.run_command(source)
else:
self.run_source(source)
except CypherError as error:
if error.classification == "ClientError":
pass
elif error.classification == "DatabaseError":
pass
elif error.classification == "TransientError":
pass
else:
pass
self.echo("{}: {}".format(error.title, error.message), err=True)
except TransactionError:
self.echo("Transaction error", err=True, fg=self.err_colour)
except ServiceUnavailable:
raise
except Exception as error:
self.echo("{}: {}".format(error.__class__.__name__, str(error)),
err=True,
fg=self.err_colour)
def begin_transaction(self):
if self.tx is None:
self.tx = self.graph.begin()
self.tx_counter = 1
self.echo(u"--- BEGIN at {} ---".format(datetime.now()),
err=True,
fg=self.tx_colour,
bold=True)
else:
self.echo(u"Transaction already open",
err=True,
fg=self.err_colour)
def commit_transaction(self):
if self.tx:
try:
self.tx.commit()
self.echo(u"--- COMMIT at {} ---".format(datetime.now()),
err=True,
fg=self.tx_colour,
bold=True)
finally:
self.tx = None
self.tx_counter = 0
else:
self.echo(u"No current transaction", err=True, fg=self.err_colour)
def rollback_transaction(self):
if self.tx:
try:
self.tx.rollback()
self.echo(u"--- ROLLBACK at {} ---".format(datetime.now()),
err=True,
fg=self.tx_colour,
bold=True)
finally:
self.tx = None
self.tx_counter = 0
else:
self.echo(u"No current transaction", err=True, fg=self.err_colour)
def read(self):
if self.multi_line:
self.multi_line = False
return self.prompt(u"", multiline=True, **self.prompt_args)
def get_prompt_tokens(_):
tokens = []
if self.tx is None:
tokens.append((Token.Prompt, "\n-> "))
else:
tokens.append((Token.Prompt, "\n-("))
tokens.append((Token.TxCounter, "{}".format(self.tx_counter)))
tokens.append((Token.Prompt, ")-> "))
return tokens
return self.prompt(get_prompt_tokens=get_prompt_tokens,
**self.prompt_args)
def run_source(self, source):
for i, statement in enumerate(self.lexer.get_statements(source)):
if i > 0:
self.echo(u"")
if statement.upper() == "BEGIN":
self.begin_transaction()
elif statement.upper() == "COMMIT":
self.commit_transaction()
elif statement.upper() == "ROLLBACK":
self.rollback_transaction()
elif self.tx is None:
self.run_cypher(self.graph.run, statement, {})
else:
self.run_cypher(self.tx.run,
statement, {},
line_no=self.tx_counter)
self.tx_counter += 1
def run_cypher(self, runner, statement, parameters, line_no=0):
t0 = timer()
result = runner(statement, parameters)
record_count = self.write_result(result)
status = u"{} record{} from {} in {:.3f}s".format(
record_count,
"" if record_count == 1 else "s",
address_str(result.summary().server.address),
timer() - t0,
)
if line_no:
self.echo(u"(", err=True, fg=self.meta_colour, bold=True, nl=False)
self.echo(u"{}".format(line_no),
err=True,
fg=self.tx_colour,
bold=True,
nl=False)
self.echo(u")->({})".format(status),
err=True,
fg=self.meta_colour,
bold=True)
else:
self.echo(u"({})".format(status),
err=True,
fg=self.meta_colour,
bold=True)
def write_result(self, result, page_size=50):
table = Table(result)
table_size = len(table)
for skip in range(0, table_size, page_size):
self.result_writer(table,
file=self.output_file,
header={
"fg": "cyan",
"bold": True
},
skip=skip,
limit=page_size)
self.echo("\r\n", nl=False)
return table_size
def run_command(self, source):
source = source.lstrip()
assert source
terms = shlex.split(source)
command_name = terms[0]
try:
command = self.commands[command_name]
except KeyError:
self.echo("Unknown command: " + command_name,
err=True,
fg=self.err_colour)
else:
args = []
kwargs = {}
for term in terms[1:]:
if "=" in term:
key, _, value = term.partition("=")
kwargs[key] = value
else:
args.append(term)
command(*args, **kwargs)
def set_multi_line(self, **kwargs):
self.multi_line = True
def edit(self, **kwargs):
initial_message = b""
with NamedTemporaryFile(suffix=".cypher") as f:
f.write(initial_message)
f.flush()
call([EDITOR, f.name])
f.seek(0)
source = f.read().decode("utf-8")
self.echo(source)
self.run(source)
def help(self, **kwargs):
self.echo(DESCRIPTION, err=True)
self.echo(u"", err=True)
self.echo(FULL_HELP.replace("\b\n", ""), err=True)
def exit(self, **kwargs):
exit(0)
def play(self, file_name):
work = self.load_unit_of_work(file_name=file_name)
with self.graph.begin() as tx:
work(tx)
def load_unit_of_work(self, file_name):
""" Load a transaction function from a cypher source file.
"""
with open(expanduser(file_name)) as f:
source = f.read()
def unit_of_work(tx):
for line_no, statement in enumerate(
self.lexer.get_statements(source), start=1):
if line_no > 0:
self.echo(u"")
self.run_cypher(tx.run, statement, {}, line_no=line_no)
return unit_of_work
def set_csv_result_writer(self, **kwargs):
self.result_writer = Table.write_csv
def set_tabular_result_writer(self, **kwargs):
self.result_writer = Table.write
def set_tsv_result_writer(self, **kwargs):
self.result_writer = Table.write_tsv
def config(self, **kwargs):
result = self.graph.run("CALL dbms.listConfig")
records = None
last_category = None
for record in result:
name = record["name"]
category, _, _ = name.partition(".")
if category != last_category:
if records is not None:
Table(records, ["name", "value"]).write(auto_align=False,
padding=0,
separator=u" = ")
self.echo(u"")
records = []
records.append((name, record["value"]))
last_category = category
if records is not None:
Table(records, ["name", "value"]).write(auto_align=False,
padding=0,
separator=u" = ")
def kernel(self, **kwargs):
result = self.graph.run(
"CALL dbms.queryJmx",
{"query": "org.neo4j:instance=kernel#0,name=Kernel"})
records = []
for record in result:
attributes = record["attributes"]
for key, value_dict in sorted(attributes.items()):
value = value_dict["value"]
if key.endswith("Date") or key.endswith("Time"):
try:
value = datetime.fromtimestamp(value /
1000).isoformat(" ")
except:
pass
records.append((key, value))
Table(records, ["key", "value"]).write(auto_align=False,
padding=0,
separator=u" = ")
class ClientConsole(Console):
multi_line = False
def __init__(self, profile=None, *_, **settings):
super(ClientConsole,
self).__init__("py2neo", verbosity=settings.get("verbosity", 0))
self.output_file = settings.pop("file", None)
welcome = settings.get("welcome", True)
if welcome:
self.write(TITLE)
self.write()
self.write(dedent(QUICK_HELP))
self.write()
self.profile = ConnectionProfile(profile, **settings)
routing = settings.get("routing", False)
try:
self.graph = Graph(self.profile,
routing=routing) # TODO: use Connector instead
except OSError as error:
self.critical("Could not connect to <%s> (%s)", self.profile.uri,
" ".join(map(str, error.args)))
raise
else:
self.debug("Connected to <%s>", self.graph.service.uri)
try:
makedirs(HISTORY_FILE_DIR)
except OSError:
pass
self.history = FileHistory(path_join(HISTORY_FILE_DIR, HISTORY_FILE))
self.lexer = CypherLexer()
self.result_writer = Table.write
self.commands = {
"//": self.set_multi_line,
"/e": self.edit,
"/?": self.help,
"/h": self.help,
"/help": self.help,
"/x": self.exit,
"/exit": self.exit,
"/play": self.play,
"/csv": self.set_csv_result_writer,
"/table": self.set_tabular_result_writer,
"/tsv": self.set_tsv_result_writer,
"/config": self.config,
"/kernel": self.kernel,
}
self.tx = None
self.qid = 0
def process_all(self, lines, times=1):
gap = False
for _ in range(times):
for line in lines:
if gap:
self.write("")
self.process(line)
if not is_command(line):
gap = True
return 0
def process(self, line):
line = line.strip()
if not line:
return
try:
if is_command(line):
self.run_command(line)
else:
self.run_source(line)
except (Neo4jError, Failure) as error:
# TODO: once this class wraps a Connector instead of a
# Graph and the errors raised by that class are only
# Failures and not Neo4jErrors, this only needs to
# catch Failure.
if hasattr(error, "title") and hasattr(error, "message"):
self.error("%s: %s", error.title, error.message)
else:
self.error("%s: %s", error.__class__.__name__,
" ".join(map(str, error.args)))
except OSError as error:
self.critical("Service Unavailable (%s)", error.args[0])
except Exception as error:
self.exception(*error.args)
def begin_transaction(self):
if self.tx is None:
self.tx = self.graph.begin()
self.qid = 1
else:
self.warning("Transaction already open")
def commit_transaction(self):
if self.tx:
try:
self.tx.commit()
self.info("Transaction committed")
finally:
self.tx = None
self.qid = 0
else:
self.warning("No current transaction")
def rollback_transaction(self):
if self.tx:
try:
self.tx.rollback()
self.info("Transaction rolled back")
finally:
self.tx = None
self.qid = 0
else:
self.warning("No current transaction")
def read(self):
prompt_args = {
"history":
self.history,
"lexer":
PygmentsLexer(CypherLexer),
"style":
merge_styles([
style_from_pygments_cls(NativeStyle),
style_from_pygments_dict({
Token.Prompt.User: "******",
Token.Prompt.At: "#ansiteal",
Token.Prompt.Host: "#ansiteal",
Token.Prompt.QID: "#ansiyellow",
Token.Prompt.Arrow: "#808080",
})
])
}
self.write()
if self.multi_line:
self.multi_line = False
return prompt(u"", multiline=True, **prompt_args)
def get_prompt_tokens():
tokens = [
("class:pygments.prompt.user", self.profile.user),
("class:pygments.prompt.at", "@"),
("class:pygments.prompt.host", self.profile.host),
]
if self.tx is None:
tokens.append(("class:pygments.prompt.arrow", " -> "))
else:
tokens.append(("class:pygments.prompt.arrow", " "))
tokens.append(("class:pygments.prompt.qid", str(self.qid)))
tokens.append(("class:pygments.prompt.arrow", "> "))
return tokens
return prompt(get_prompt_tokens, **prompt_args)
def run_source(self, source):
for i, statement in enumerate(self.lexer.get_statements(source)):
if i > 0:
self.write(u"")
if statement.upper() == "BEGIN":
self.begin_transaction()
elif statement.upper() == "COMMIT":
self.commit_transaction()
elif statement.upper() == "ROLLBACK":
self.rollback_transaction()
elif self.tx is None:
self.run_cypher(self.graph.run, statement, {})
else:
self.run_cypher(self.tx.run, statement, {}, query_id=self.qid)
self.qid += 1
def run_cypher(self, runner, statement, parameters, query_id=0):
t0 = timer()
result = runner(statement, parameters)
record_count = self.write_result(result)
summary = result.summary()
if summary.connection:
uri = summary.connection["uri"]
else:
uri = self.graph.service.uri
msg = "Fetched %r %s from %r in %rs"
args = [
record_count,
"record" if record_count == 1 else "records",
uri,
timer() - t0,
]
if query_id:
msg += " for query (%r)"
args.append(query_id)
self.debug(msg, *args)
def write_result(self, result, page_size=50):
table = Table(result)
table_size = len(table)
if self.verbosity >= 0:
for skip in range(0, table_size, page_size):
self.result_writer(table,
file=self.output_file,
header="cyan",
skip=skip,
limit=page_size)
self.write("\r\n", end='')
return table_size
def run_command(self, source):
source = source.lstrip()
assert source
terms = shlex.split(source)
command_name = terms[0]
try:
command = self.commands[command_name]
except KeyError:
self.info("Unknown command: " + command_name)
else:
args = []
kwargs = {}
for term in terms[1:]:
if "=" in term:
key, _, value = term.partition("=")
kwargs[key] = value
else:
args.append(term)
command(*args, **kwargs)
def set_multi_line(self, **kwargs):
self.multi_line = True
def edit(self, **kwargs):
initial_message = b""
with NamedTemporaryFile(suffix=".cypher") as f:
f.write(initial_message)
f.flush()
call([EDITOR, f.name])
f.seek(0)
source = f.read().decode("utf-8")
self.write(source)
self.process(source)
def help(self, **kwargs):
self.info(DESCRIPTION)
self.info(u"")
self.info(FULL_HELP.replace("\b\n", ""))
def play(self, file_name):
work = self.load_unit_of_work(file_name=file_name)
with self.graph.begin() as tx:
work(tx)
def load_unit_of_work(self, file_name):
""" Load a transaction function from a cypher source file.
"""
with open(expanduser(file_name)) as f:
source = f.read()
def unit_of_work(tx):
for line_no, statement in enumerate(
self.lexer.get_statements(source), start=1):
if line_no > 0:
self.write(u"")
self.run_cypher(tx.run, statement, {}, query_id=line_no)
return unit_of_work
def set_csv_result_writer(self, **kwargs):
self.result_writer = Table.write_csv
def set_tabular_result_writer(self, **kwargs):
self.result_writer = Table.write
def set_tsv_result_writer(self, **kwargs):
self.result_writer = Table.write_tsv
def config(self, **kwargs):
result = self.graph.run("CALL dbms.listConfig")
records = None
last_category = None
for record in result:
name = record["name"]
category, _, _ = name.partition(".")
if category != last_category:
if records is not None:
Table(records, ["name", "value"]).write(auto_align=False,
padding=0,
separator=u" = ")
self.write(u"")
records = []
records.append((name, record["value"]))
last_category = category
if records is not None:
Table(records, ["name", "value"]).write(auto_align=False,
padding=0,
separator=u" = ")
def kernel(self, **kwargs):
result = self.graph.run(
"CALL dbms.queryJmx",
{"query": "org.neo4j:instance=kernel#0,name=Kernel"})
records = []
for record in result:
attributes = record["attributes"]
for key, value_dict in sorted(attributes.items()):
value = value_dict["value"]
if key.endswith("Date") or key.endswith("Time"):
try:
value = datetime.fromtimestamp(value /
1000).isoformat(" ")
except:
pass
records.append((key, value))
Table(records, ["key", "value"]).write(auto_align=False,
padding=0,
separator=u" = ")
to_node = Node('manufacturer',
manuf_name=to_node_name[0],
manuf_comment=to_node_name[1])
if None in [from_node, to_node]:
raise Exception('Unknown connection type', connection)
rel = Relationship(from_node, connection, to_node)
graph.merge(rel)
def worker(mac_address):
print(mac_address)
manuf_data = mac_parser.get_all(mac_address)
if manuf_data is None or manuf_data.manuf is None:
return
register_connection('ETHER/GEN/OUI/MANUFACTURED_BY', mac_address,
(manuf_data.manuf, manuf_data.comment))
graph = Graph(
password='******') # TODO: parameterize, don't hardcode password
mac_parser = manuf.MacParser(update=True)
if __name__ == '__main__':
pool = multiprocessing.Pool(multiprocessing.cpu_count())
manuf.MacParser(update=True)
macs = [
_['d.mac_address'] for _ in graph.run(
'MATCH (d:device) WHERE NOT (d)-[:`ETHER/GEN/OUI/MANUFACTURED_BY`]-() RETURN d.mac_address'
)
]
pool.map(worker, macs)