Files
WGDashboard/src/modules/AmneziaWireguardConfiguration.py
2025-08-20 15:06:27 +08:00

372 lines
19 KiB
Python

"""
AmneziaWG Configuration
"""
import random, sqlalchemy, os, subprocess, re, uuid
from .PeerJobs import PeerJobs
from .AmneziaWGPeer import AmneziaWGPeer
from .PeerShareLinks import PeerShareLinks
from .Utilities import RegexMatch
from .WireguardConfiguration import WireguardConfiguration
class AmneziaWireguardConfiguration(WireguardConfiguration):
def __init__(self, DashboardConfig,
AllPeerJobs: PeerJobs,
AllPeerShareLinks: PeerShareLinks,
name: str = None, data: dict = None, backup: dict = None, startup: bool = False):
self.Jc = 0
self.Jmin = 0
self.Jmax = 0
self.S1 = 0
self.S2 = 0
self.H1 = 1
self.H2 = 2
self.H3 = 3
self.H4 = 4
super().__init__(DashboardConfig, AllPeerJobs, AllPeerShareLinks, name, data, backup, startup, wg=False)
def toJson(self):
self.Status = self.getStatus()
return {
"Status": self.Status,
"Name": self.Name,
"PrivateKey": self.PrivateKey,
"PublicKey": self.PublicKey,
"Address": self.Address,
"ListenPort": self.ListenPort,
"PreUp": self.PreUp,
"PreDown": self.PreDown,
"PostUp": self.PostUp,
"PostDown": self.PostDown,
"SaveConfig": self.SaveConfig,
"DataUsage": {
"Total": sum(list(map(lambda x: x.cumu_data + x.total_data, self.Peers))),
"Sent": sum(list(map(lambda x: x.cumu_sent + x.total_sent, self.Peers))),
"Receive": sum(list(map(lambda x: x.cumu_receive + x.total_receive, self.Peers)))
},
"ConnectedPeers": len(list(filter(lambda x: x.status == "running", self.Peers))),
"TotalPeers": len(self.Peers),
"Protocol": self.Protocol,
"Jc": self.Jc,
"Jmin": self.Jmin,
"Jmax": self.Jmax,
"S1": self.S1,
"S2": self.S2,
"H1": self.H1,
"H2": self.H2,
"H3": self.H3,
"H4": self.H4
}
def createDatabase(self, dbName = None):
if dbName is None:
dbName = self.Name
self.peersTable = sqlalchemy.Table(
dbName, self.metadata,
sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False, primary_key=True),
sqlalchemy.Column('private_key', sqlalchemy.String(255)),
sqlalchemy.Column('DNS', sqlalchemy.Text),
sqlalchemy.Column('advanced_security', sqlalchemy.String(255)),
sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
sqlalchemy.Column('name', sqlalchemy.Text),
sqlalchemy.Column('total_receive', sqlalchemy.Float),
sqlalchemy.Column('total_sent', sqlalchemy.Float),
sqlalchemy.Column('total_data', sqlalchemy.Float),
sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
sqlalchemy.Column('status', sqlalchemy.String(255)),
sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
sqlalchemy.Column('cumu_data', sqlalchemy.Float),
sqlalchemy.Column('mtu', sqlalchemy.Integer),
sqlalchemy.Column('keepalive', sqlalchemy.Integer),
sqlalchemy.Column('remote_endpoint', sqlalchemy.String(255)),
sqlalchemy.Column('preshared_key', sqlalchemy.String(255)),
extend_existing=True
)
self.peersRestrictedTable = sqlalchemy.Table(
f'{dbName}_restrict_access', self.metadata,
sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False, primary_key=True),
sqlalchemy.Column('private_key', sqlalchemy.String(255)),
sqlalchemy.Column('DNS', sqlalchemy.Text),
sqlalchemy.Column('advanced_security', sqlalchemy.String(255)),
sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
sqlalchemy.Column('name', sqlalchemy.Text),
sqlalchemy.Column('total_receive', sqlalchemy.Float),
sqlalchemy.Column('total_sent', sqlalchemy.Float),
sqlalchemy.Column('total_data', sqlalchemy.Float),
sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
sqlalchemy.Column('status', sqlalchemy.String(255)),
sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
sqlalchemy.Column('cumu_data', sqlalchemy.Float),
sqlalchemy.Column('mtu', sqlalchemy.Integer),
sqlalchemy.Column('keepalive', sqlalchemy.Integer),
sqlalchemy.Column('remote_endpoint', sqlalchemy.String(255)),
sqlalchemy.Column('preshared_key', sqlalchemy.String(255)),
extend_existing=True
)
self.peersTransferTable = sqlalchemy.Table(
f'{dbName}_transfer', self.metadata,
sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False),
sqlalchemy.Column('total_receive', sqlalchemy.Float),
sqlalchemy.Column('total_sent', sqlalchemy.Float),
sqlalchemy.Column('total_data', sqlalchemy.Float),
sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
sqlalchemy.Column('cumu_data', sqlalchemy.Float),
sqlalchemy.Column('time', (sqlalchemy.DATETIME if self.DashboardConfig.GetConfig("Database", "type")[1] == 'sqlite' else sqlalchemy.TIMESTAMP),
server_default=sqlalchemy.func.now()),
extend_existing=True
)
self.peersDeletedTable = sqlalchemy.Table(
f'{dbName}_deleted', self.metadata,
sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False),
sqlalchemy.Column('private_key', sqlalchemy.String(255)),
sqlalchemy.Column('DNS', sqlalchemy.Text),
sqlalchemy.Column('advanced_security', sqlalchemy.String(255)),
sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
sqlalchemy.Column('name', sqlalchemy.Text),
sqlalchemy.Column('total_receive', sqlalchemy.Float),
sqlalchemy.Column('total_sent', sqlalchemy.Float),
sqlalchemy.Column('total_data', sqlalchemy.Float),
sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
sqlalchemy.Column('status', sqlalchemy.String(255)),
sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
sqlalchemy.Column('cumu_data', sqlalchemy.Float),
sqlalchemy.Column('mtu', sqlalchemy.Integer),
sqlalchemy.Column('keepalive', sqlalchemy.Integer),
sqlalchemy.Column('remote_endpoint', sqlalchemy.String(255)),
sqlalchemy.Column('preshared_key', sqlalchemy.String(255)),
extend_existing=True
)
self.metadata.create_all(self.engine)
# existingTables = sqlSelect("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
# existingTables = [t['name'] for t in existingTables]
# if dbName not in existingTables:
# sqlUpdate(
# """
# CREATE TABLE '%s'(
# id VARCHAR NOT NULL, private_key VARCHAR NULL, DNS VARCHAR NULL, advanced_security VARCHAR NULL,
# endpoint_allowed_ip VARCHAR NULL, name VARCHAR NULL, total_receive FLOAT NULL,
# total_sent FLOAT NULL, total_data FLOAT NULL, endpoint VARCHAR NULL,
# status VARCHAR NULL, latest_handshake VARCHAR NULL, allowed_ip VARCHAR NULL,
# cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, mtu INT NULL,
# keepalive INT NULL, remote_endpoint VARCHAR NULL, preshared_key VARCHAR NULL,
# PRIMARY KEY (id)
# )
# """ % dbName
# )
#
# if f'{dbName}_restrict_access' not in existingTables:
# sqlUpdate(
# """
# CREATE TABLE '%s_restrict_access' (
# id VARCHAR NOT NULL, private_key VARCHAR NULL, DNS VARCHAR NULL, advanced_security VARCHAR NULL,
# endpoint_allowed_ip VARCHAR NULL, name VARCHAR NULL, total_receive FLOAT NULL,
# total_sent FLOAT NULL, total_data FLOAT NULL, endpoint VARCHAR NULL,
# status VARCHAR NULL, latest_handshake VARCHAR NULL, allowed_ip VARCHAR NULL,
# cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, mtu INT NULL,
# keepalive INT NULL, remote_endpoint VARCHAR NULL, preshared_key VARCHAR NULL,
# PRIMARY KEY (id)
# )
# """ % dbName
# )
# if f'{dbName}_transfer' not in existingTables:
# sqlUpdate(
# """
# CREATE TABLE '%s_transfer' (
# id VARCHAR NOT NULL, total_receive FLOAT NULL,
# total_sent FLOAT NULL, total_data FLOAT NULL,
# cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, time DATETIME
# )
# """ % dbName
# )
# if f'{dbName}_deleted' not in existingTables:
# sqlUpdate(
# """
# CREATE TABLE '%s_deleted' (
# id VARCHAR NOT NULL, private_key VARCHAR NULL, DNS VARCHAR NULL, advanced_security VARCHAR NULL,
# endpoint_allowed_ip VARCHAR NULL, name VARCHAR NULL, total_receive FLOAT NULL,
# total_sent FLOAT NULL, total_data FLOAT NULL, endpoint VARCHAR NULL,
# status VARCHAR NULL, latest_handshake VARCHAR NULL, allowed_ip VARCHAR NULL,
# cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, mtu INT NULL,
# keepalive INT NULL, remote_endpoint VARCHAR NULL, preshared_key VARCHAR NULL,
# PRIMARY KEY (id)
# )
# """ % dbName
# )
def getPeers(self):
self.Peers.clear()
if self.configurationFileChanged():
with open(self.configPath, 'r') as configFile:
p = []
pCounter = -1
content = configFile.read().split('\n')
try:
peerStarts = content.index("[Peer]")
content = content[peerStarts:]
for i in content:
if not RegexMatch("#(.*)", i) and not RegexMatch(";(.*)", i):
if i == "[Peer]":
pCounter += 1
p.append({})
p[pCounter]["name"] = ""
else:
if len(i) > 0:
split = re.split(r'\s*=\s*', i, 1)
if len(split) == 2:
p[pCounter][split[0]] = split[1]
if RegexMatch("#Name# = (.*)", i):
split = re.split(r'\s*=\s*', i, 1)
if len(split) == 2:
p[pCounter]["name"] = split[1]
with self.engine.begin() as conn:
for i in p:
if "PublicKey" in i.keys():
tempPeer = conn.execute(self.peersTable.select().where(
self.peersTable.columns.id == i['PublicKey']
)).mappings().fetchone()
if tempPeer is None:
tempPeer = {
"id": i['PublicKey'],
"advanced_security": i.get('AdvancedSecurity', 'off'),
"private_key": "",
"DNS": self.DashboardConfig.GetConfig("Peers", "peer_global_DNS")[1],
"endpoint_allowed_ip": self.DashboardConfig.GetConfig("Peers", "peer_endpoint_allowed_ip")[
1],
"name": i.get("name"),
"total_receive": 0,
"total_sent": 0,
"total_data": 0,
"endpoint": "N/A",
"status": "stopped",
"latest_handshake": "N/A",
"allowed_ip": i.get("AllowedIPs", "N/A"),
"cumu_receive": 0,
"cumu_sent": 0,
"cumu_data": 0,
"mtu": self.DashboardConfig.GetConfig("Peers", "peer_mtu")[1],
"keepalive": self.DashboardConfig.GetConfig("Peers", "peer_keep_alive")[1],
"remote_endpoint": self.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1],
"preshared_key": i["PresharedKey"] if "PresharedKey" in i.keys() else ""
}
# sqlUpdate(
# """
# INSERT INTO '%s'
# VALUES (:id, :private_key, :DNS, :advanced_security, :endpoint_allowed_ip, :name, :total_receive, :total_sent,
# :total_data, :endpoint, :status, :latest_handshake, :allowed_ip, :cumu_receive, :cumu_sent,
# :cumu_data, :mtu, :keepalive, :remote_endpoint, :preshared_key);
# """ % self.Name
# , newPeer)
conn.execute(
self.peersTable.insert().values(tempPeer)
)
else:
# sqlUpdate("UPDATE '%s' SET allowed_ip = ? WHERE id = ?" % self.Name,
# (i.get("AllowedIPs", "N/A"), i['PublicKey'],))
conn.execute(
self.peersTable.update().values({
"allowed_ip": i.get("AllowedIPs", "N/A")
}).where(
self.peersTable.columns.id == i['PublicKey']
)
)
self.Peers.append(AmneziaWGPeer(tempPeer, self))
except Exception as e:
if __name__ == '__main__':
print(f"[WGDashboard] {self.Name} getPeers() Error: {str(e)}")
else:
# checkIfExist = sqlSelect("SELECT * FROM '%s'" % self.Name).fetchall()
with self.engine.connect() as conn:
existingPeers = conn.execute(self.peersTable.select()).mappings().fetchall()
for i in existingPeers:
self.Peers.append(AmneziaWGPeer(i, self))
def addPeers(self, peers: list) -> tuple[bool, dict]:
result = {
"message": None,
"peers": []
}
try:
with self.engine.begin() as conn:
for i in peers:
newPeer = {
"id": i['id'],
"private_key": i['private_key'],
"DNS": i['DNS'],
"endpoint_allowed_ip": i['endpoint_allowed_ip'],
"name": i['name'],
"total_receive": 0,
"total_sent": 0,
"total_data": 0,
"endpoint": "N/A",
"status": "stopped",
"latest_handshake": "N/A",
"allowed_ip": i.get("allowed_ip", "N/A"),
"cumu_receive": 0,
"cumu_sent": 0,
"cumu_data": 0,
"mtu": i['mtu'],
"keepalive": i['keepalive'],
"remote_endpoint": self.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1],
"preshared_key": i["preshared_key"],
"advanced_security": i['advanced_security']
}
conn.execute(
self.peersTable.insert().values(newPeer)
)
# sqlUpdate(
# """
# INSERT INTO '%s'
# VALUES (:id, :private_key, :DNS, :advanced_security, :endpoint_allowed_ip, :name, :total_receive, :total_sent,
# :total_data, :endpoint, :status, :latest_handshake, :allowed_ip, :cumu_receive, :cumu_sent,
# :cumu_data, :mtu, :keepalive, :remote_endpoint, :preshared_key);
# """ % self.Name
# , newPeer)
for p in peers:
presharedKeyExist = len(p['preshared_key']) > 0
rd = random.Random()
uid = str(uuid.UUID(int=rd.getrandbits(128), version=4))
if presharedKeyExist:
with open(uid, "w+") as f:
f.write(p['preshared_key'])
subprocess.check_output(
f"{self.Protocol} set {self.Name} peer {p['id']} allowed-ips {p['allowed_ip'].replace(' ', '')}{f' preshared-key {uid}' if presharedKeyExist else ''}",
shell=True, stderr=subprocess.STDOUT)
if presharedKeyExist:
os.remove(uid)
subprocess.check_output(
f"{self.Protocol}-quick save {self.Name}", shell=True, stderr=subprocess.STDOUT)
self.getPeersList()
for p in peers:
p = self.searchPeer(p['id'])
if p[0]:
result['peers'].append(p[1])
return True, result
except Exception as e:
result['message'] = str(e)
return False, result
def getRestrictedPeers(self):
self.RestrictedPeers = []
# restricted = sqlSelect("SELECT * FROM '%s_restrict_access'" % self.Name).fetchall()
with self.engine.connect() as conn:
restricted = conn.execute(self.peersRestrictedTable.select()).mappings().fetchall()
for i in restricted:
self.RestrictedPeers.append(AmneziaWGPeer(i, self))