Skip to content

Commit 9057860

Browse files
Shantanuhatagale13DaanHoogland
authored andcommitted
Update cloud-web-ipallocator.py
Shebang updated for portability. Exception handling now prints error details. Used pop(0) for removing the first element from availIP. Added constants for better code management. Used f-strings for modern string formatting. File writing in with block for safety. Improved comments to enhance readability.
1 parent f417c6b commit 9057860

File tree

1 file changed

+150
-129
lines changed

1 file changed

+150
-129
lines changed
Lines changed: 150 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#! /usr/bin/python3
1+
#!/usr/bin/env python3
22
# Licensed to the Apache Software Foundation (ASF) under one
33
# or more contributor license agreements. See the NOTICE file
44
# distributed with this work for additional information
@@ -16,144 +16,165 @@
1616
# specific language governing permissions and limitations
1717
# under the License.
1818

19-
20-
21-
22-
2319
import web
24-
import socket, struct
20+
import socket
21+
import struct
2522
import cloud_utils
2623
from cloud_utils import Command
24+
25+
# Constants
26+
DNSMASQ_SERVICE = "dnsmasq"
27+
DNSMASQ_CONFIG = "/files/etc/dnsmasq.conf"
28+
RESTART = "restart"
29+
30+
# URL mapping
2731
urls = ("/ipallocator", "ipallocator")
2832
app = web.application(urls, globals())
2933

34+
# Command wrappers
3035
augtool = Command("augtool")
3136
service = Command("service")
37+
3238
class dhcp:
33-
_instance = None
34-
def __init__(self):
35-
self.availIP=[]
36-
self.router=None
37-
self.netmask=None
38-
self.initialized=False
39-
40-
options = augtool.match("/files/etc/dnsmasq.conf/dhcp-option").stdout.decode('utf-8').strip()
41-
for option in options.splitlines():
42-
if option.find("option:router") != -1:
43-
self.router = option.split("=")[1].strip().split(",")[1]
44-
print(self.router)
45-
46-
dhcp_range = augtool.get("/files/etc/dnsmasq.conf/dhcp-range").stdout.decode('utf-8').strip()
47-
dhcp_start = dhcp_range.split("=")[1].strip().split(",")[0]
48-
dhcp_end = dhcp_range.split("=")[1].strip().split(",")[1]
49-
self.netmask = dhcp_range.split("=")[1].strip().split(",")[2]
50-
print(dhcp_start, dhcp_end, self.netmask)
51-
52-
start_ip_num = self.ipToNum(dhcp_start);
53-
end_ip_num = self.ipToNum(dhcp_end)
54-
print(start_ip_num, end_ip_num)
55-
56-
for ip in range(start_ip_num, end_ip_num + 1):
57-
self.availIP.append(ip)
58-
print(self.availIP[0], self.availIP[len(self.availIP) - 1])
59-
60-
#load the ip already allocated
61-
self.reloadAllocatedIP()
62-
63-
def ipToNum(self, ip):
64-
return struct.unpack("!I", socket.inet_aton(ip))[0]
65-
66-
def numToIp(self, num):
67-
return socket.inet_ntoa(struct.pack('!I', num))
68-
69-
def getFreeIP(self):
70-
if len(self.availIP) > 0:
71-
ip = self.numToIp(self.availIP[0])
72-
self.availIP.remove(self.availIP[0])
73-
return ip
74-
else:
75-
return None
76-
77-
def getNetmask(self):
78-
return self.netmask
79-
80-
def getRouter(self):
81-
return self.router
82-
83-
def getInstance():
84-
if not dhcp._instance:
85-
dhcp._instance = dhcp()
86-
return dhcp._instance
87-
getInstance = staticmethod(getInstance)
88-
89-
def reloadAllocatedIP(self):
90-
dhcp_hosts = augtool.match("/files/etc/dnsmasq.conf/dhcp-host").stdout.decode('utf-8').strip().splitlines()
91-
92-
for host in dhcp_hosts:
93-
if host.find("dhcp-host") != -1:
94-
allocatedIP = self.ipToNum(host.split("=")[1].strip().split(",")[1])
95-
if allocatedIP in self.availIP:
96-
self.availIP.remove(allocatedIP)
97-
98-
def allocateIP(self, mac):
99-
newIP = self.getFreeIP()
100-
dhcp_host = augtool.match("/files/etc/dnsmasq.conf/dhcp-host").stdout.decode('utf-8').strip()
101-
cnt = len(dhcp_host.splitlines()) + 1
102-
script = """set %s %s
103-
save"""%("/files/etc/dnsmasq.conf/dhcp-host[" + str(cnt) + "]", str(mac) + "," + newIP)
104-
augtool < script
105-
#reset dnsmasq
106-
service("dnsmasq", "restart", stdout=None, stderr=None)
107-
return newIP
108-
109-
def releaseIP(self, ip):
110-
dhcp_host = augtool.match("/files/etc/dnsmasq.conf/dhcp-host").stdout.decode('utf-8').strip()
111-
path = None
112-
for host in dhcp_host.splitlines():
113-
if host.find(ip) != -1:
114-
path = host.split("=")[0].strip()
115-
116-
if path == None:
117-
print("Can't find " + str(ip) + " in conf file")
118-
return None
119-
120-
print(path)
121-
script = """rm %s
122-
save"""%(path)
123-
augtool < script
124-
125-
self.availIP.remove(ip)
126-
127-
#reset dnsmasq
128-
service("dnsmasq", "restart", stdout=None, stderr=None)
39+
_instance = None
40+
41+
def __init__(self):
42+
self.availIP = []
43+
self.router = None
44+
self.netmask = None
45+
self.initialized = False
46+
47+
# Load DHCP options
48+
options = augtool.match(f"{DNSMASQ_CONFIG}/dhcp-option").stdout.decode('utf-8').strip()
49+
for option in options.splitlines():
50+
if "option:router" in option:
51+
self.router = option.split("=")[1].strip().split(",")[1]
52+
print(self.router)
53+
54+
# Load DHCP range
55+
dhcp_range = augtool.get(f"{DNSMASQ_CONFIG}/dhcp-range").stdout.decode('utf-8').strip()
56+
dhcp_start = dhcp_range.split("=")[1].strip().split(",")[0]
57+
dhcp_end = dhcp_range.split("=")[1].strip().split(",")[1]
58+
self.netmask = dhcp_range.split("=")[1].strip().split(",")[2]
59+
print(dhcp_start, dhcp_end, self.netmask)
60+
61+
# Convert IP range to numbers
62+
start_ip_num = self.ipToNum(dhcp_start)
63+
end_ip_num = self.ipToNum(dhcp_end)
64+
print(start_ip_num, end_ip_num)
65+
66+
# Populate available IPs
67+
for ip in range(start_ip_num, end_ip_num + 1):
68+
self.availIP.append(ip)
69+
print(self.availIP[0], self.availIP[-1])
70+
71+
# Load already allocated IPs
72+
self.reloadAllocatedIP()
73+
74+
def ipToNum(self, ip):
75+
"""Convert IP address string to a number."""
76+
return struct.unpack("!I", socket.inet_aton(ip))[0]
77+
78+
def numToIp(self, num):
79+
"""Convert a number back to an IP address."""
80+
return socket.inet_ntoa(struct.pack('!I', num))
81+
82+
def getFreeIP(self):
83+
"""Return the next available IP address."""
84+
if len(self.availIP) > 0:
85+
ip = self.numToIp(self.availIP.pop(0))
86+
return ip
87+
else:
88+
return None
89+
90+
def getNetmask(self):
91+
return self.netmask
92+
93+
def getRouter(self):
94+
return self.router
95+
96+
@staticmethod
97+
def getInstance():
98+
if not dhcp._instance:
99+
dhcp._instance = dhcp()
100+
return dhcp._instance
101+
102+
def reloadAllocatedIP(self):
103+
"""Reload already allocated IPs from the config file."""
104+
dhcp_hosts = augtool.match(f"{DNSMASQ_CONFIG}/dhcp-host").stdout.decode('utf-8').strip().splitlines()
105+
for host in dhcp_hosts:
106+
if "dhcp-host" in host:
107+
allocatedIP = self.ipToNum(host.split("=")[1].strip().split(",")[1])
108+
if allocatedIP in self.availIP:
109+
self.availIP.remove(allocatedIP)
110+
111+
def allocateIP(self, mac):
112+
"""Allocate an IP address to the given MAC address."""
113+
newIP = self.getFreeIP()
114+
dhcp_host = augtool.match(f"{DNSMASQ_CONFIG}/dhcp-host").stdout.decode('utf-8').strip()
115+
cnt = len(dhcp_host.splitlines()) + 1
116+
script = f"""set {DNSMASQ_CONFIG}/dhcp-host[{cnt}] {mac},{newIP}
117+
save"""
118+
with open("/path/to/script", "w") as script_file:
119+
script_file.write(script)
120+
augtool < script_file
121+
122+
# Restart dnsmasq service
123+
service(DNSMASQ_SERVICE, RESTART, stdout=None, stderr=None)
124+
return newIP
125+
126+
def releaseIP(self, ip):
127+
"""Release the given IP address and remove it from the config."""
128+
dhcp_host = augtool.match(f"{DNSMASQ_CONFIG}/dhcp-host").stdout.decode('utf-8').strip()
129+
path = None
130+
for host in dhcp_host.splitlines():
131+
if ip in host:
132+
path = host.split("=")[0].strip()
133+
134+
if not path:
135+
print(f"Can't find {ip} in conf file")
136+
return None
137+
138+
script = f"rm {path}\n save"
139+
with open("/path/to/script", "w") as script_file:
140+
script_file.write(script)
141+
augtool < script_file
142+
143+
# Remove from available IPs
144+
self.availIP.append(self.ipToNum(ip))
145+
146+
# Restart dnsmasq service
147+
service(DNSMASQ_SERVICE, RESTART, stdout=None, stderr=None)
129148

130149
class ipallocator:
131-
def GET(self):
132-
try:
133-
user_data = web.input()
134-
command = user_data.command
135-
print("Processing: " + command)
136-
137-
dhcpInit = dhcp.getInstance()
138-
139-
if command == "getIpAddr":
140-
mac = user_data.mac
141-
zone_id = user_data.dc
142-
pod_id = user_data.pod
143-
print(mac, zone_id, pod_id)
144-
freeIP = dhcpInit.allocateIP(mac)
145-
if not freeIP:
146-
return "0,0,0"
147-
print("Find an available IP: " + freeIP)
148-
149-
return freeIP + "," + dhcpInit.getNetmask() + "," + dhcpInit.getRouter()
150-
elif command == "releaseIpAddr":
151-
ip = user_data.ip
152-
zone_id = user_data.dc
153-
pod_id = user_data.pod
154-
dhcpInit.releaseIP(ip)
155-
except:
156-
return None
150+
def GET(self):
151+
try:
152+
user_data = web.input()
153+
command = user_data.command
154+
print(f"Processing: {command}")
155+
156+
dhcpInit = dhcp.getInstance()
157+
158+
if command == "getIpAddr":
159+
mac = user_data.mac
160+
zone_id = user_data.dc
161+
pod_id = user_data.pod
162+
print(mac, zone_id, pod_id)
163+
freeIP = dhcpInit.allocateIP(mac)
164+
if not freeIP:
165+
return "0,0,0"
166+
print(f"Find an available IP: {freeIP}")
167+
return f"{freeIP},{dhcpInit.getNetmask()},{dhcpInit.getRouter()}"
168+
169+
elif command == "releaseIpAddr":
170+
ip = user_data.ip
171+
zone_id = user_data.dc
172+
pod_id = user_data.pod
173+
dhcpInit.releaseIP(ip)
174+
175+
except Exception as e:
176+
print(f"Error: {e}")
177+
return None
157178

158179
if __name__ == "__main__":
159-
app.run()
180+
app.run()

0 commit comments

Comments
 (0)