PicoBadUSB/code.py at main · hackeringtrue/PicoBadUSB · GitHub

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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import wifi
import socketpool
import ipaddress
import sys
import gc
import time
import os

from atoms.networking import WebHost
from atoms.config_man import ConfigMan
from atoms.hid import AtomDucky, load_payload_from_file
from atoms.buttons import pixel
from atoms.colors import color

pixel.fill(color("yellow"))

class AP:
    def __init__(self, ssid="Atom Ducky", passw="", ap='TRUE', ip="10.0.0.15", mode="NORMAL"):
        self.ssid = ssid
        self.password = passw
        self.ip = ip
        self.ap = ap
        self.mode = mode
        self.config = ConfigMan()
        self.run()

    def parse_config(self, config_data):
        for line in config_data.splitlines():
            key, value = line.split('=')
            if key == 'IP':
                self.ip = value
            elif key == 'SSID':
                self.ssid = value
            elif key == 'PASSW':
                self.password = value
            elif key == 'MODE':
                self.mode = value
            elif key == 'AP':
                self.ap = value

    def run(self):
        try:
            contents = os.listdir('/atoms')
            cfg_exists = '_config' in contents
            if cfg_exists:
                print("Config exists! Reading the _config..")
                _config = self.config.read_config()
                self.parse_config(_config)
                print(f"""IP: {self.ip}
SSID: {self.ssid}
PASSW: {len(self.password) * '*'}
MODE: {self.mode}
AP: {self.ap}
""")
        except OSError:
            print("You probably need to enable read/write permissions, copy Atom Ducky's boot.py to your drive and press RESET!")
        time.sleep(0.2)
        self.rubber_mode_exec()

        if self.ap == 'FALSE':
            self.connect_to_ap()
        else:
            self.config_ips()
            self.start_host_ap()

        if not cfg_exists:
            self.config = ConfigMan(ip=self.ip, ssid=self.ssid, passw=self.password, mode=self.mode, ap=self.ap)
            self.config.write_config()

        self.config.edit_config(ip=self.ip)
        print(f"Changed the _config IP to: {self.ip}")

        webserver = WebHost(ip=self.ip, port=80)

    def rubber_mode_exec(self):
        if self.mode == "RUBBER":
            pixel.fill(color("red"))
            print("Detected RUBBER mode! Injecting first...")
            ducky = AtomDucky()
            payload = load_payload_from_file()
            if payload is not None:
                ducky.payloads_write(payload)
            gc.collect()
            pixel.fill(color("green"))
            time.sleep(0.5)

    def connect_to_ap(self):
        try:
            wifi.radio.hostname = "Atom-Ducky"
            print("Connecting to {}...".format(self.ssid))
            wifi.radio.connect(self.ssid, self.password)
            while not wifi.radio.connected:
                time.sleep(0.5)
            self.ip = wifi.radio.ipv4_address
            print("Connected! IP Address of the Atom Ducky: {}".format(self.ip))
        except ConnectionError as e:
            print("Connection failed! Failswitching to AP mode... Error reason:", str(e))
            self.ip = "10.0.0.15"
            self.config_ips()
            self.start_host_ap()

    def start_host_ap(self):
        wifi.radio.start_ap(ssid=self.ssid, password=self.password)
        print("Access Point started with SSID:", self.ssid)

    def config_ips(self):
        if not self.ip:
            self.ip = "10.0.0.15"
        ip = ipaddress.IPv4Address(self.ip)
        netmask = ipaddress.IPv4Address("255.255.255.0")
        gateway = ipaddress.IPv4Address(self.ip)
        wifi.radio.set_ipv4_address_ap(ipv4=ip, netmask=netmask, gateway=gateway)

def shutdown():
    ip = ipaddress.IPv4Address('192.168.1.42')
    netmask = ipaddress.IPv4Address("255.255.255.255")
    gateway = ipaddress.IPv4Address('192.168.1.42')
    wifi.radio.set_ipv4_address_ap(ipv4=ip, netmask=netmask, gateway=gateway)
    wifi.radio.stop_ap()
    wifi.radio.enabled = False

def main():
    try:
        ap = AP()
        time.sleep(2)

    except KeyboardInterrupt:
        print("Catched KeyboardInterrupt")
        shutdown()

if __name__ == "__main__":
    main()