import pathlib import urllib.error import urllib.request from typing import Union, Iterable, Dict, Any, List from dataclasses import dataclass from .general import SysCommand from .output import info, warn from .exceptions import SysCallError from .storage import storage @dataclass class CustomMirror: url: str signcheck: str signoptions: str name: str def sort_mirrorlist(raw_data :bytes, sort_order: List[str] = ['https', 'http']) -> bytes: """ This function can sort /etc/pacman.d/mirrorlist according to the mirror's URL prefix. By default places HTTPS before HTTP but it also preserves the country/rank-order. This assumes /etc/pacman.d/mirrorlist looks like the following: ## Comment Server = url or ## Comment #Server = url But the Comments need to start with double-hashmarks to be distringuished from server url definitions (commented or uncommented). """ comments_and_whitespaces = b"" sort_order += ['Unknown'] categories: Dict[str, List] = {key: [] for key in sort_order} for line in raw_data.split(b"\n"): if line[0:2] in (b'##', b''): comments_and_whitespaces += line + b'\n' elif line[:6].lower() == b'server' or line[:7].lower() == b'#server': opening, url = line.split(b'=', 1) opening, url = opening.strip(), url.strip() if (category := url.split(b'://',1)[0].decode('UTF-8')) in categories: categories[category].append(comments_and_whitespaces) categories[category].append(opening + b' = ' + url + b'\n') else: categories["Unknown"].append(comments_and_whitespaces) categories["Unknown"].append(opening + b' = ' + url + b'\n') comments_and_whitespaces = b"" new_raw_data = b'' for category in sort_order + ["Unknown"]: for line in categories[category]: new_raw_data += line return new_raw_data def filter_mirrors_by_region(regions :str, destination :str = '/etc/pacman.d/mirrorlist', sort_order :List[str] = ["https", "http"], *args :str, **kwargs :str ) -> Union[bool, bytes]: """ This function will change the active mirrors on the live medium by filtering which regions are active based on `regions`. :param regions: A series of country codes separated by `,`. For instance `SE,US` for sweden and United States. :type regions: str """ region_list = [f'country={region}' for region in regions.split(',')] response = urllib.request.urlopen(urllib.request.Request(f"https://archlinux.org/mirrorlist/?{'&'.join(region_list)}&protocol=https&protocol=http&ip_version=4&ip_version=6&use_mirror_status=on'", headers={'User-Agent': 'ArchInstall'})) new_list = response.read().replace(b"#Server", b"Server") if sort_order: new_list = sort_mirrorlist(new_list, sort_order=sort_order) if destination: with open(destination, "wb") as mirrorlist: mirrorlist.write(new_list) return True else: return new_list.decode('UTF-8') def add_custom_mirrors(mirrors: List[CustomMirror]) -> bool: """ This will append custom mirror definitions in pacman.conf :param mirrors: A list of custom mirrors :type mirrors: List[CustomMirror] """ with open('/etc/pacman.conf', 'a') as pacman: for mirror in mirrors: pacman.write(f"[{mirror.name}]\n") pacman.write(f"SigLevel = {mirror.signcheck} {mirror.signoptions}\n") pacman.write(f"Server = {mirror.url}\n") return True def insert_mirrors(mirrors :Dict[str, Any], *args :str, **kwargs :str) -> bool: """ This function will insert a given mirror-list at the top of `/etc/pacman.d/mirrorlist`. It will not flush any other mirrors, just insert new ones. :param mirrors: A dictionary of `{'url' : 'country', 'url2' : 'country'}` :type mirrors: dict """ original_mirrorlist = '' with open('/etc/pacman.d/mirrorlist', 'r') as original: original_mirrorlist = original.read() with open('/etc/pacman.d/mirrorlist', 'w') as new_mirrorlist: for mirror, country in mirrors.items(): new_mirrorlist.write(f'## {country}\n') new_mirrorlist.write(f'Server = {mirror}\n') new_mirrorlist.write('\n') new_mirrorlist.write(original_mirrorlist) return True def use_mirrors( regions: Dict[str, Iterable[str]], destination: str = '/etc/pacman.d/mirrorlist' ): info(f'A new package mirror-list has been created: {destination}') with open(destination, 'w') as mirrorlist: for region, mirrors in regions.items(): for mirror in mirrors: mirrorlist.write(f'## {region}\n') mirrorlist.write(f'Server = {mirror}\n') def re_rank_mirrors( top: int = 10, src: str = '/etc/pacman.d/mirrorlist', dst: str = '/etc/pacman.d/mirrorlist', ) -> bool: try: cmd = SysCommand(f"/usr/bin/rankmirrors -n {top} {src}") except SysCallError: return False with open(dst, 'w') as f: f.write(str(cmd)) return True def list_mirrors(sort_order :List[str] = ["https", "http"]) -> Dict[str, Any]: regions: Dict[str, Dict[str, Any]] = {} if storage['arguments']['offline']: with pathlib.Path('/etc/pacman.d/mirrorlist').open('rb') as fh: mirrorlist = fh.read() else: url = "https://archlinux.org/mirrorlist/?protocol=https&protocol=http&ip_version=4&ip_version=6&use_mirror_status=on" try: response = urllib.request.urlopen(url) except urllib.error.URLError as err: warn(f'Could not fetch an active mirror-list: {err}') return regions mirrorlist = response.read() if sort_order: mirrorlist = sort_mirrorlist(mirrorlist, sort_order=sort_order) region = 'Unknown region' for line in mirrorlist.split(b'\n'): if len(line.strip()) == 0: continue clean_line = line.decode('UTF-8').strip('\n').strip('\r') if clean_line[:3] == '## ': region = clean_line[3:] elif clean_line[:10] == '#Server = ': regions.setdefault(region, {}) url = clean_line.lstrip('#Server = ') regions[region][url] = True elif clean_line.startswith('Server = '): regions.setdefault(region, {}) url = clean_line.lstrip('Server = ') regions[region][url] = True return regions