-
Notifications
You must be signed in to change notification settings - Fork 1
/
cli.py
79 lines (56 loc) · 1.7 KB
/
cli.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import logging
import os
import click
from lib.ui_base import BaseUI
from lib.utils import TexttableWithLogStream, LogStream
if "LOG_FILE" in os.environ:
log_file = os.environ["LOG_FILE"]
else:
log_file = "iot-device-id.log"
log_stream = LogStream()
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(threadName)-12.12s] [%(levelname)-5.5s] %(message)s",
handlers=[
logging.FileHandler(log_file),
logging.StreamHandler(log_stream)
])
class ConsoleUI(BaseUI):
def __init__(self):
super().__init__()
self.tt = None
self.rows = []
self.sort_by = 0
def set_headers(self, headers):
self.tt = TexttableWithLogStream(log_stream, headers)
def sort_by_row(self, i):
self.sort_by = i
def add_row(self, values):
self.rows.append(values)
def draw(self):
rows = self.rows
self.rows = []
for v in rows:
self.tt.add_row(v)
self.tt.draw()
def record(self, ip, name):
"""Record a new device into the database."""
self.start_recording(ip, name)
def detect(self):
"""Detect IoT devices on the local network."""
self.start_detecting()
@click.group()
def cli():
pass
@click.command()
@click.option("--ip", help="IP address of target device", prompt="Target IP address")
@click.option("--name", prompt="Device name", help="Name of the target device")
def record(ip, name):
ConsoleUI().record(ip, name)
@click.command()
def detect():
ConsoleUI().detect()
cli.add_command(record)
cli.add_command(detect)
if __name__ == '__main__':
cli()