pylcg/pylcg.py
2024-11-25 22:28:06 -05:00

134 lines
4.3 KiB
Python

#!/usr/bin/env python3
# Python implementation of a Linear Congruential Generator for IP Sharding - Developed by acidvegas in Python (https://git.acid.vegas/pylcg)
# pylcg.py
import argparse
import asyncio
import ipaddress
from math import ceil
class LCG:
'''Linear Congruential Generator for deterministic random number generation'''
def __init__(self, seed: int, m: int = 2**32):
self.m = m
self.a = 1597
self.c = 51749
self.seed = seed
self.current = seed
def get_nth(self, n: int) -> int:
'''
Get the nth number in the sequence without generating previous numbers.
:param n: The index of the number to get
'''
# For large n, use the standard next() method to avoid modular arithmetic issues
if n > 1000:
self.current = self.seed
for _ in range(n):
self.next()
return self.current
# For smaller n, use direct calculation
result = self.seed
for _ in range(n):
result = (self.a * result + self.c) % self.m
return result
def next(self) -> int:
'''Generate next random number'''
self.current = (self.a * self.current + self.c) % self.m
return self.current
class IPRange:
'''Memory-efficient IP range iterator'''
def __init__(self, cidr: str):
network = ipaddress.ip_network(cidr)
self.start = int(network.network_address)
self.end = int(network.broadcast_address)
self.total = self.end - self.start + 1
def get_ip_at_index(self, index: int) -> str:
'''
Get IP at specific index without generating previous IPs
:param index: The index of the IP to get
'''
if not 0 <= index < self.total:
raise IndexError('IP index out of range')
return str(ipaddress.ip_address(self.start + index))
async def get_shard_ips(cidr: str, shard_num: int, total_shards: int, seed: int, chunk_size: int = 1000):
'''
Asynchronously generate IPs for the specified shard.
:param cidr: The CIDR range to shard
:param shard_num: The number of the shard to generate
:param total_shards: The total number of shards
:param seed: The seed for the random number generator
:param chunk_size: The size of the chunks to process
'''
# Initialize the IP range and LCG
ip_range = IPRange(cidr)
lcg = LCG(seed)
total_ips = ip_range.total
# Calculate which indices belong to this shard
shard_size = ceil(total_ips / total_shards)
start_idx = shard_num * shard_size
end_idx = min(start_idx + shard_size, total_ips)
# Process in chunks to maintain memory efficiency
for chunk_start in range(start_idx, end_idx, chunk_size):
chunk_end = min(chunk_start + chunk_size, end_idx)
chunk_indices = list(range(chunk_start, chunk_end))
# Generate random values for this chunk
chunk_random_values = [(i, lcg.get_nth(i)) for i in chunk_indices]
chunk_random_values.sort(key=lambda x: x[1])
# Yield IPs in randomized order
for idx, _ in chunk_random_values:
yield ip_range.get_ip_at_index(idx)
# Allow other tasks to run (do we need this?)
await asyncio.sleep(0)
async def main():
parser = argparse.ArgumentParser(description='Async IP address sharding tool')
parser.add_argument('cidr', help='Target IP range in CIDR format')
parser.add_argument('shard_num', type=int, help='Shard number (0-based)')
parser.add_argument('total_shards', type=int, help='Total number of shards')
parser.add_argument('--seed', type=int, default=12345, help='Random seed for LCG')
parser.add_argument('--chunk-size', type=int, default=1000, help='Processing chunk size')
args = parser.parse_args()
if args.shard_num >= args.total_shards:
raise ValueError('Shard number must be less than total shards')
if args.shard_num < 0 or args.total_shards < 1:
raise ValueError('Invalid shard configuration')
async for ip in get_shard_ips(args.cidr, args.shard_num, args.total_shards, args.seed, args.chunk_size):
print(ip)
if __name__ == '__main__':
asyncio.run(main())