Skip to content

Commit b1e47df

Browse files
WillowSauceRWillowSauceR
authored andcommitted
Optimize code, use multiprocess to receive packets
1 parent 911ace1 commit b1e47df

File tree

1 file changed

+62
-51
lines changed

1 file changed

+62
-51
lines changed

scan.py

Lines changed: 62 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,24 @@
1+
from multiprocessing.connection import _ConnectionBase
12
import socket
23
import sys
34
import threading
45
import time
56
import os
67
import re
78
from random import randint
9+
import multiprocessing as mp
10+
from webbrowser import get
811

9-
from api import getLocalHostIP, log
12+
from api import getLocalHostIP, getTime, log
1013

1114
localHostIP = getLocalHostIP()
1215
localHostPort = randint(1024, 65535)
1316

1417
socketSendRecv = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
1518
socketSendRecv.bind((localHostIP, localHostPort))
1619

20+
scanResult = {"serverCount": 0, "bdsCount": 0, "nkCount": 0, "geyserCount": 0, "skipped": 0, "error": 0, "serverList": [], "totalPlayerCount": 0}
21+
1722
try:
1823
TargetAddr = sys.argv[1]
1924
portRange = sys.argv[2]
@@ -34,7 +39,6 @@
3439
fileName = ""
3540

3641
motdData = b'\x01\x00\x00\x00\x00$\r\x12\xd3\x00\xff\xff\x00\xfe\xfe\xfe\xfe\xfd\xfd\xfd\xfd\x124Vx\n'
37-
serverCount = 0
3842

3943

4044
def getIpList(ip: str):
@@ -78,11 +82,12 @@ def sendPacket(portStart, portEnd, dstAddr):
7882
if stopThread:
7983
break
8084
if dstPort % int(portEnd / 5) == 0 and verboseMode == "y":
81-
log(f"Scaning port: {str(dstPort)} ~ {str(dstPort + int(portEnd / 5))}", info = "I")
85+
log(f"Scaning port: {str(dstPort)} ~ {str(dstPort + int(portEnd / 5))}", info="I")
8286
socketSendRecv.sendto(motdData, (dstAddr, dstPort))
8387
# if verboseMode == "y":
8488
# log(f"Port {portStart} ~ {portEnd} Done", info = "I")
8589

90+
8691
def startThreads():
8792
global stopThread
8893
if "-" in portRange:
@@ -95,57 +100,52 @@ def startThreads():
95100
for dstAddr in ipList:
96101
stopThread = False
97102
print()
98-
log(f"Scaning target: {dstAddr}", info = "I")
103+
log(f"Scaning target: {dstAddr}", info="I")
99104
print()
100105
# sendPacket(portStart, portEnd, dstAddr)
101-
t1 = threading.Thread(target=sendPacket, args=(portStart, portEnd, dstAddr), daemon = True)
106+
t1 = threading.Thread(target=sendPacket, args=(
107+
portStart, portEnd, dstAddr), daemon=True)
102108
t1.start()
103-
tmpServerCount = serverCount
109+
tmpServerCount = scanResult['serverCount']
104110
if timeout != 0:
105111
time.sleep(timeout)
106-
if tmpServerCount == serverCount:
112+
if tmpServerCount == scanResult['serverCount']:
107113
stopThread = True
108114
t1.join()
109115
# while threading.enumerate().__len__() != 2: # main and itself
110116
# time.sleep(1)
111117

112-
log("BE Server Count: " + str(serverCount), info = "I")
113-
log("BDS Count: " + str(bdsCount), info = "I")
114-
log("NK Count: " + str(nkCount), info = "I")
115-
log("Geyser Count: " + str(geyserCount), info = "I")
116-
log("Skipped Count: " + str(skipped), info = "I")
117-
log("Error Count: " + str(error), info = "I")
118-
log("Total Player Count: " + str(totalPlayerCount), info = "I")
118+
log(f"BE Server Count: {scanResult['serverCount']}", info="I")
119+
log(f"BDS Count: {scanResult['bdsCount']}", info="I")
120+
log(f"NK Count: {scanResult['nkCount']}", info="I")
121+
log(f"Geyser Count: {scanResult['geyserCount']}", info="I")
122+
log(f"Skipped Count: {scanResult['skipped']}", info="I")
123+
log(f"Error Count: {scanResult['error']}", info="I")
124+
log(f"Total Player Count: {scanResult['totalPlayerCount']}", info="I")
119125
os._exit(0)
120126

121-
if __name__ == "__main__":
122-
t = threading.Thread(target=startThreads, daemon = True)
123-
t.start()
124127

125-
bdsCount = 0
126-
nkCount = 0
127-
geyserCount = 0
128-
skipped = 0
129-
error = 0
130-
payloads = []
131-
totalPlayerCount = 0
128+
def recvPackets(socketSendRecv: socket.socket, verboseMode: str, fileName: str, scanResult: dict, pipe: _ConnectionBase):
132129
while True:
133130
try:
134131
data, addr = socketSendRecv.recvfrom(10240)
135-
date = time.strftime('%H:%M:%S')
132+
date = getTime()
136133
if len(data) <= 30:
137134
skipped += 1
138-
log(f"data length <= 30, may not motd packet, skipped. Source: {addr[0]}:{addr[1]}", info = "R")
135+
log(
136+
f"data length <= 30, may not motd packet, skipped. Source: {addr[0]}:{addr[1]}", info="R")
139137
continue
140138
if b"MCPE" not in data:
141139
skipped += 1
142-
log(f"metadata \"MCPE\" not in packet, may not motd packet, skipped. Source: {addr[0]}:{addr[1]}", info = "R")
140+
log(
141+
f"metadata \"MCPE\" not in packet, may not motd packet, skipped. Source: {addr[0]}:{addr[1]}", info="R")
143142
continue
144-
if addr not in payloads:
145-
payloads.append(addr)
143+
if addr not in scanResult['serverList']:
144+
scanResult['serverList'].append(addr)
146145
else:
147146
skipped += 1
148-
log(f"Duplicate server found, skipped. Source: {addr[0]}:{addr[1]}", info = "R")
147+
log(
148+
f"Duplicate server found, skipped. Source: {addr[0]}:{addr[1]}", info="R")
149149
continue
150150
infos = []
151151
data1 = data.split(b"MCPE")
@@ -160,41 +160,52 @@ def startThreads():
160160
context = str(info)[2:-1]
161161
infos.append(context)
162162
# log(len(infos))
163-
log(f"Motd: {infos[1]}", info = "R")
164-
log(f"Versin: {infos[3]}/{infos[2]}", info = "R")
165-
log(f"Online: {infos[4]}/{infos[5]}", info = "R")
163+
log(f"Motd: {infos[1]}", info="R")
164+
log(f"Versin: {infos[3]}/{infos[2]}", info="R")
165+
log(f"Online: {infos[4]}/{infos[5]}", info="R")
166166
try:
167-
log(f"Map: {infos[7]}/{infos[8]}", info = "R")
167+
log(f"Map: {infos[7]}/{infos[8]}", info="R")
168168
except:
169-
log(f"Map info is unavailable.", info = "R")
169+
log(f"Map info is unavailable.", info="R")
170170
try:
171-
log(f"Port(v4/v6): {infos[10]}/{infos[11]}", info = "R")
171+
log(f"Port(v4/v6): {infos[10]}/{infos[11]}", info="R")
172172
except:
173-
log(f"Port info is unavailable.", info = "R")
174-
log(f"Source: {addr[0]}:{addr[1]}", info = "R")
175-
serverCount += 1
176-
totalPlayerCount += int(infos[4])
177-
log(f"{str(serverCount)}", info = "C")
178-
log(f"{totalPlayerCount}", info = "P")
173+
log(f"Port info is unavailable.", info="R")
174+
log(f"Source: {addr[0]}:{addr[1]}", info="R")
175+
scanResult['serverCount'] += 1
176+
scanResult['totalPlayerCount'] += int(infos[4])
177+
log(f"{scanResult['serverCount']}", info="C")
178+
log(f"{scanResult['totalPlayerCount']}", info="P")
179179
print("")
180180
if fileName:
181181
with open(fileName, "a+") as file:
182182
file.write(
183-
f"{date} | {serverCount} | {addr[0]} | {addr[1]} | {infos[1]} | {infos[3]} | {infos[4]} | {infos[5]}\n")
183+
f"{date} | {scanResult['serverCount']} | {addr[0]} | {addr[1]} | {infos[1]} | {infos[3]} | {infos[4]} | {infos[5]}\n")
184184
if len(infos) == 10 or len(infos) == 6:
185185
nkCount += 1
186186
elif re.search(b"edicated", data):
187-
bdsCount += 1
187+
scanResult['bdsCount'] += 1
188188
elif re.search(b"nukkit", data):
189-
nkCount += 1
189+
scanResult['nkCount'] += 1
190190
elif re.search(b"eyser", data):
191-
geyserCount += 1
192-
#socketSendRecv.close()
191+
scanResult['geyserCount'] += 1
192+
# socketSendRecv.close()
193+
pipe.send(scanResult)
193194
except OSError as errorInfo:
194195
if verboseMode == "y":
195-
log(f"{errorInfo}, skipped.", info = "R")
196-
error += 1
196+
log(f"{errorInfo}, skipped.", info="R")
197+
scanResult['error'] += 1
197198
except Exception as errorInfo:
198-
log(f"{errorInfo}, skipped.", info = "R")
199-
error += 1
199+
log(f"{errorInfo}, skipped.", info="R")
200+
scanResult['error'] += 1
200201
pass
202+
203+
204+
if __name__ == "__main__":
205+
t = threading.Thread(target=startThreads, daemon=True)
206+
t.start()
207+
pipe1, pipe2 = mp.Pipe()
208+
p = mp.Process(target=recvPackets, args=(socketSendRecv, verboseMode, fileName, scanResult, pipe2), daemon=True)
209+
p.start()
210+
while mp.active_children():
211+
scanResult = pipe2.recv()

0 commit comments

Comments
 (0)