diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 35564d6..de3eab2 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -42,12 +42,20 @@ We have implemented GET and HEAD, as there is no requirement for any other metho The HEAD method is used by some PXE ROMs to find the Content-Length before the GET is sent. +## NBD +NBD is similar to NFS in that it can act as a root device for Linux systems. Defined in the [specification](https://github.com/yoe/nbd/blob/master/doc/proto.txt), NBD allows access to block devices over the network by performing read and write requests on the block device itself. + +This is different to NFS as it does not act as a filesystem, merely a single file. NBD supports read/write access along with copy-on-write support, both in memory and on disk. Read/write without copy-on-write is potentially dangerous if the file-system layer does not support multiple systems accessing it at the same time. Copy-on-write alleviates these potential problems by providing a volatile layer in which per-client changes are saved. Both the on-disk and in-memory configurations delete changes after the client disconnects, but the in-memory configuration may offer a speed increase as the changes are stored in the system RAM. + +WARNING: The use of this option can potentially consume a large amount of RAM; up to the size of the disk image multiplied by the number of connected clients may be used. The same can be said for the on-disk configuration, where this configuration uses disk space rather than memory space. A further configuration option to store the original disk image in memory is provided to potentially allow for read/write speed up. + # PyPXE Services -The PyPXE library provies the following services for the purpose of creating a Python-based PXE environment: TFTP, HTTP, and DHCP. Each service must be imorted independently as such: +The PyPXE library provides the following services for the purpose of creating a Python-based PXE environment: TFTP, HTTP, DHCP, and NBD. Each service must be imported independently as such: * `from pypxe import tftp` or `import pypxe.tftp` imports the TFTP service * `from pypxe import dhcp` or `import pypxe.dhcp` imports the DHCP service * `from pypxe import http` or `import pypxe.http` imports the HTTP service +* `from pypxe import nbd` or `import pypxe.nbd` imports the NBD service **See [`pypxe-server.py`](pypxe-server.py) in the root of the repo for example usage on how to call, define, and setup the services.** When running any Python script that uses these classes, it should be run as a user with root privileges as they bind to interfaces and without root privileges the services will most likely fail to bind properly. @@ -70,7 +78,7 @@ The TFTP server class, __`TFTPD()`__, is constructed with the following __keywor |__`ip`__|This is the IP address that the TFTP server will bind to.|`'0.0.0.0'` (so that it binds to all available interfaces)| _string_| |__`port`__|This it the port that the TFTP server will run on.|`69` (default port for TFTP)|_int_| |__`netboot_directory`__|This is the directory that the TFTP server will serve files from similarly to that of `tftpboot`.|`'.'` (current directory)|_string_| -|__`mode_debug`__|This indicates whether or not the TFTP server should be started in debug mode or not.|`False`|`bool`| +|__`mode_debug`__|This indicates whether or not the TFTP server should be started in debug mode or not.|`False`|_bool_| |__`logger`__|A [Logger](https://docs.python.org/2/library/logging.html#logger-objects) object used for logging messages, if `None` a local [StreamHandler](https://docs.python.org/2/library/logging.handlers.html#streamhandler) instance will be created.|`None`|[_Logger_](https://docs.python.org/2/library/logging.html#logger-objects)| |__`default_retries`__|The number of data retransmissions before dropping a connection.|`3`|_int_| |__`timeout`__|The time in seconds before re-sending an un-acknowledged data block.|`5`|_int_| @@ -86,7 +94,7 @@ from pypxe import dhcp import pypxe.dhcp ``` -###Usage +### Usage The DHCP server class, __`DHCPD()`__, is constructed with the following __keyword arguments__: |Keyword Argument|Description|Default|Type| @@ -104,6 +112,8 @@ The DHCP server class, __`DHCPD()`__, is constructed with the following __keywor |__`use_ipxe`__|This indicates whether or not iPXE is being used and adjusts itself accordingly.|`False`|_bool_| |__`use_http`__|This indicates whether or not the built-in HTTP server is being used and adjusts itself accordingly.|`False`|_bool_| |__`mode_proxy`__|This indicates whether or not the DHCP server should be started in ProxyDHCP mode or not.|`False`|_bool_| +|__`static_config`__|This specifies a static configuration dictionary so that it can give specific leases to specific MAC addresses.|`{}`|_dict_| +|__`whitelist`__|This indicates whether or not the DHCP server should use the static configuration dictionary as a whitelist; effectively, the DHCP server will only give out leases to those specified in the `static_config` dictionary.|`False`|_bool_| |__`mode_debug`__|This indicates whether or not the DHCP server should be started in debug mode or not.|`False`|_bool_| |__`logger`__|A [Logger](https://docs.python.org/2/library/logging.html#logger-objects) object used for logging messages, if `None` a local [StreamHandler](https://docs.python.org/2/library/logging.handlers.html#streamhandler) instance will be created.|`None`|[_Logger_](https://docs.python.org/2/library/logging.html#logger-objects)| @@ -126,9 +136,36 @@ The HTTP server class, __`HTTPD()`__, is constructed with the following __keywor |__`ip`__|This is the IP address that the HTTP server will bind to.|`'0.0.0.0'` (so that it binds to all available interfaces)|_string_| |__`port`__|This it the port that the HTTP server will run on.|`80` (default port for HTTP)|_int_| |__`netboot_directory`__|This is the directory that the HTTP server will serve files from similarly to that of `tftpboot`.|`'.'` (current directory)|_string_| -|__`mode_debug`__|This indicates whether or not the HTTP server should be started in debug mode or not.|`False`|bool| +|__`mode_debug`__|This indicates whether or not the HTTP server should be started in debug mode or not.|`False`|_bool_| |__`logger`__|A [Logger](https://docs.python.org/2/library/logging.html#logger-objects) object used for logging messages, if `None` a local [StreamHandler](https://docs.python.org/2/library/logging.handlers.html#streamhandler) instance will be created.|`None`|[_Logger_](https://docs.python.org/2/library/logging.html#logger-objects)| +## NBD Server `pypxe.nbd` + +### Importing +The NBD service can be imported _one_ of the following two ways: +```python +from pypxe import http +``` +```python +import pypxe.nbd +``` + +### Usage +The NBD server class, __`NBD()`__, is constructed with the following __keyword arguments__: + +|Keyword Argument|Description|Default|Type| +|---|---|---|---| +|__`ip`__|This is the IP address that the NBD server will bind to.|`'0.0.0.0'` (so that it binds to all available interfaces)|_string_| +|__`port`__|This it the port that the NBD server will run on.|`10809` (default port for NBD)|_int_| +|__`block_device`__|The filename of the block device to be used as the root device.|`''`|_string_| +|__`write`__|Enable write support on the block device.|`False`|_bool_| +|__`cow`__|Enable copy-on-write support on the block device.|`True`|_bool_| +|__`in_mem`__|Enable _in-memory_ copy-on-write support on the block device. `False` causes changes to be stored on disk.|`False`|_bool_| +|__`copy_to_ram`__|Copy the disk image to RAM when the service starts.|`False`|_bool_| +|__`mode_debug`__|This indicates whether or not the NBD server should be started in debug mode or not.|`False`|_bool_| +|__`logger`__|A [Logger](https://docs.python.org/2/library/logging.html#logger-objects) object used for logging messages, if `None` a local [StreamHandler](https://docs.python.org/2/library/logging.handlers.html#streamhandler) instance will be created.|`None`|[_Logger_](https://docs.python.org/2/library/logging.html#logger-objects)| + + ## Additional Information * The function `chr(0)` is used in multiple places throughout the servers. This denotes a `NULL` byte, or `\x00` * Python 2.6 does not include the `argparse` module, it is included in the standard library as of 2.7 and newer. The `argparse` module is required to take in command line arguments and `pypxe-server.py` will not run without it. diff --git a/README.md b/README.md index ff1fd23..0d2bf5a 100644 --- a/README.md +++ b/README.md @@ -43,12 +43,17 @@ The following are arguments that can be passed to `pypxe-server.py` when running |Argument|Description|Default| |---|---|---| |__`--ipxe`__|Enable iPXE ROM|`False`| +|__`--no-ipxe`__|Disable iPXE ROM|`True`| |__`--http`__|Enable built-in HTTP server|`False`| +|__`--no-http`__|Disable built-in HTTP server|`True`| |__`--dhcp`__|Enable built-in DHCP server|`False`| |__`--dhcp-proxy`__|Enable built-in DHCP server in proxy mode (implies `--dhcp`)|`False`| +|__`--no-dhcp`__|Disable built-in DHCP server|`True`| +|__`--tftp`__|Enable built-in TFTP server which is enabled by default|`True`| |__`--no-tftp`__|Disable built-in TFTP server which is enabled by default|`False`| -|__`--debug`__|Enable selected services in DEBUG mode; services are selected by passing the name in a comma separated list. **Options are: http, tftp and dhcp** _This adds a level of verbosity so that you can see what's happening in the background._|`''`| +|__`--debug`__|Enable selected services in DEBUG mode; services are selected by passing the name in a comma separated list. **Options are: http, tftp and dhcp**; one can also prefix an option with `-` to prevent debugging of that service; for example, the following will enable debugging for all services _except_ the DHCP service `--debug all,-dhcp`. _This mode adds a level of verbosity so that you can see what's happening in the background._|`''`| |__`--config`__|Load configuration from JSON file. (see [`example_cfg.json`](example_cfg.json))|`None`| +|__`--static-config`__|Load DHCP lease configuration from JSON file. (see [`example-leases.json`](example-leases.json))|`None`| |__`--syslog`__|Specify a syslog server|`None`| |__`--syslog-port`__|Specify a syslog server port|`514`| @@ -57,25 +62,39 @@ The following are arguments that can be passed to `pypxe-server.py` when running |Argument|Description|Default| |---|---|---| -|__`-s DHCP_SERVER_IP`__ or __`--dhcp-server-ip DHCP_SERVER_IP`__|Specify DHCP server IP address|`192.168.2.2`| -|__`-p DHCP_SERVER_PORT`__ or __`--dhcp-server-port DHCP_SERVER_PORT`__|Specify DHCP server port|`67`| -|__`-b DHCP_OFFER_BEGIN`__ or __`--dhcp-begin DHCP_OFFER_BEGIN`__|Specify DHCP lease range start|`192.168.2.100`| -|__`-e DHCP_OFFER_END`__ or __`--dhcp-end DHCP_OFFER_END`__|Specify DHCP lease range end|`192.168.2.150`| -|__`-n DHCP_SUBNET`__ or __`--dhcp-subnet DHCP_SUBNET`__|Specify DHCP subnet mask|`255.255.255.0`| -|__`-r DHCP_ROUTER`__ or __`--dhcp-router DHCP_ROUTER`__|Specify DHCP lease router|`192.168.2.1`| -|__`-d DHCP_DNS`__ or __`--dhcp-dns DHCP_DNS`__|Specify DHCP lease DNS server|`8.8.8.8`| -|__`-c DHCP_BROADCAST`__ or __`--dhcp-broadcast DHCP_BROADCAST`__|Specify DHCP broadcast address|`''`| -|__`-f DHCP_FILESERVER_IP`__ or __`--dhcp-fileserver-ip DHCP_FILESERVER_IP`__|Specify DHCP file server IP address|`192.168.2.2`| +|__`--dhcp-server-ip DHCP_SERVER_IP`__|Specify DHCP server IP address|`192.168.2.2`| +|__`--dhcp-server-port DHCP_SERVER_PORT`__|Specify DHCP server port|`67`| +|__`--dhcp-begin DHCP_OFFER_BEGIN`__|Specify DHCP lease range start|`192.168.2.100`| +|__`--dhcp-end DHCP_OFFER_END`__|Specify DHCP lease range end|`192.168.2.150`| +|__`--dhcp-subnet DHCP_SUBNET`__|Specify DHCP subnet mask|`255.255.255.0`| +| __`--dhcp-router DHCP_ROUTER`__|Specify DHCP lease router|`192.168.2.1`| +|__`--dhcp-dns DHCP_DNS`__|Specify DHCP lease DNS server|`8.8.8.8`| +|__`--dhcp-broadcast DHCP_BROADCAST`__|Specify DHCP broadcast address|`''`| +|__`--dhcp-fileserver-ip DHCP_FILESERVER_IP`__|Specify DHCP file server IP address|`192.168.2.2`| +|__`--dhcp-whitelist`__|Only serve clients specified in the static lease file (`--static-config`)|`False`| ##### File Name/Directory Arguments |Argument|Description|Default| |---|---|---| -|__`-a NETBOOT_DIR`__ or __`--netboot-dir NETBOOT_DIR`__|Specify the local directory where network boot files will be served|`'netboot'`| -|__`-i NETBOOT_FILE`__ or __`--netboot-file NETBOOT_FILE`__|Specify the PXE boot file name|_automatically set based on what services are enabled or disabled, see [`DOCUMENTATION.md`](DOCUMENTATION.md) for further explanation_| +|__`--netboot-dir NETBOOT_DIR`__|Specify the local directory where network boot files will be served|`'netboot'`| +|__`--netboot-file NETBOOT_FILE`__|Specify the PXE boot file name|_automatically set based on what services are enabled or disabled, see [`DOCUMENTATION.md`](DOCUMENTATION.md) for further explanation_| + + +##### Network Block Device Arguments +|Argument|Description|Default| +|---|---|---| +|__`--nbd NBD_BLOCK_DEVICE`__|Specify the block device to be served by NBD and enable NBD. This can be a disk image.|`''`| +|__`--nbd-write`__|Open the block device for write access. UNSAFE: Multiple clients can cause corruption|`False`| +|__`--nbd-cow`__|When write is enabled, create a *volatile* file per client with their changes. Clients can write but changes are not shared or kept.|`True (Only applies if write is on)`| +|__`--nbd-cow-in-mem`__|Client volatile changes are stored in RAM rather than on disk. WARNING: High RAM usage (up to sizeof(block device)*clients)|`False`| +|__`--nbd-copy-to-ram`__|Disk image is copied to RAM on start to speed up access. Changes are lost when write is used without cow.|`False`| +|__`--nbd-server`__|The NBD server IP address to bind to|`0.0.0.0`| +|__`--nbd-port`__|The NBD server port to bind to|`10809`| + ## Notes * `Core.iso` located in `netboot` is from the [TinyCore Project](http://distro.ibiblio.org/tinycorelinux/) and is provided as an example to network boot from using PyPXE -* `chainload.kpxe` located in `netboot` is the `undionly.kpxe` from the [iPXE Project](http://ipxe.org/) +* `chainload.kpxe` located in `netboot` is the `undionly.kpxe` from the [iPXE Project](http://ipxe.org/) * `ldlinux.c32`, `libutil.c32`, `pxelinux.0`, `menu.c32`, and `memdisk` located in `netboot` are from the [SYSLINUX Project](http://www.syslinux.org/) version [6.02](http://www.syslinux.org/wiki/index.php/Syslinux_6_Changelog#Changes_in_6.02) diff --git a/example_cfg.json b/example_cfg.json index 98cb347..7bf99c9 100644 --- a/example_cfg.json +++ b/example_cfg.json @@ -16,4 +16,11 @@ "USE_TFTP" : true, "USE_DHCP" : true, "DHCP_MODE_PROXY" : false, + "NBD_BLOCK_DEVICE" : "", + "NBD_WRITE" : false, + "NBD_COW" : true, + "NBD_COW_IN_MEM" : false, + "NBD_COPY_TO_RAM" : false, + "NBD_SERVER_IP" : "0.0.0.0", + "NBD_SERVER_PORT" : 10809, "MODE_DEBUG" : "" } diff --git a/example_leases.json b/example_leases.json new file mode 100644 index 0000000..c07150c --- /dev/null +++ b/example_leases.json @@ -0,0 +1,15 @@ +{ + "dhcp": { + "binding": { + "5C:9A:D8:5E:F2:E5": { + "dns": [ + "8.8.8.8", + "8.8.4.4" + ], + "ipaddr": "192.168.0.123", + "router": "192.168.0.1", + "subnet": "255.255.255.0" + } + } + } +} diff --git a/netboot/boot.http.nbd.ipxe b/netboot/boot.http.nbd.ipxe new file mode 100644 index 0000000..16ef087 --- /dev/null +++ b/netboot/boot.http.nbd.ipxe @@ -0,0 +1,4 @@ +#!ipxe +initrd http://192.168.0.11/initramfs-linux.img +chain http://192.168.0.11/vmlinuz-linux quiet ip=::::::dhcp nbd_host=192.168.0.11 nbd_name=arch.img root=/dev/nbd0 +sanboot diff --git a/pypxe-server.py b/pypxe-server.py index 066feba..f3c6e5c 100755 --- a/pypxe-server.py +++ b/pypxe-server.py @@ -15,6 +15,7 @@ from pypxe import tftp # PyPXE TFTP service from pypxe import dhcp # PyPXE DHCP service from pypxe import http # PyPXE HTTP service +from pypxe import nbd # PyPXE NBD service # default settings SETTINGS = {'NETBOOT_DIR':'netboot', @@ -35,39 +36,77 @@ 'USE_TFTP':True, 'USE_DHCP':True, 'DHCP_MODE_PROXY':False, + 'NBD_BLOCK_DEVICE':'', + 'NBD_WRITE':False, + 'NBD_COW':True, + 'NBD_COW_IN_MEM':False, + 'NBD_COPY_TO_RAM':False, + 'NBD_SERVER_IP':'0.0.0.0', + 'NBD_PORT':10809, 'MODE_DEBUG':''} def parse_cli_arguments(): # main service arguments parser = argparse.ArgumentParser(description = 'Set options at runtime. Defaults are in %(prog)s', formatter_class = argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--ipxe', action = 'store_true', dest = 'USE_IPXE', help = 'Enable iPXE ROM', default = SETTINGS['USE_IPXE']) - parser.add_argument('--http', action = 'store_true', dest = 'USE_HTTP', help = 'Enable built-in HTTP server', default = SETTINGS['USE_HTTP']) - parser.add_argument('--no-tftp', action = 'store_false', dest = 'USE_TFTP', help = 'Disable built-in TFTP server, by default it is enabled', default = SETTINGS['USE_TFTP']) - parser.add_argument('--debug', action = 'store', dest = 'MODE_DEBUG', help = 'Comma Seperated (http,tftp,dhcp). Adds verbosity to the selected services while they run. Use \'all\' for enabling debug on all services', default = SETTINGS['MODE_DEBUG']) + + ipxeexclusive = parser.add_mutually_exclusive_group(required = False) + ipxeexclusive.add_argument('--ipxe', action = 'store_true', dest = 'USE_IPXE', help = 'Enable iPXE ROM', default = SETTINGS['USE_IPXE']) + ipxeexclusive.add_argument('--no-ipxe', action = 'store_false', dest = 'USE_IPXE', help = 'Disable iPXE ROM', default = not SETTINGS['USE_IPXE']) + + httpexclusive = parser.add_mutually_exclusive_group(required = False) + httpexclusive.add_argument('--http', action = 'store_true', dest = 'USE_HTTP', help = 'Disable built-in HTTP server', default = SETTINGS['USE_HTTP']) + httpexclusive.add_argument('--no-http', action = 'store_false', dest = 'USE_HTTP', help = 'Enable built-in HTTP server', default = not SETTINGS['USE_HTTP']) + + tftpexclusive = parser.add_mutually_exclusive_group(required = False) + tftpexclusive.add_argument('--tftp', action = 'store_true', dest = 'USE_TFTP', help = 'Enable built-in TFTP server, by default it is enabled', default = SETTINGS['USE_TFTP']) + tftpexclusive.add_argument('--no-tftp', action = 'store_false', dest = 'USE_TFTP', help = 'Disable built-in TFTP server, by default it is enabled', default = not SETTINGS['USE_TFTP']) + + parser.add_argument('--debug', action = 'store', dest = 'MODE_DEBUG', help = 'Comma Seperated (http,tftp,dhcp). Adds verbosity to the selected services while they run. Use \'all\' for enabling debug on all services. Precede an option with \'-\' to disable debugging for that service; as an example, one can pass in the following to enable debugging for all services except the DHCP service: \'--debug all,-dhcp\'', default = SETTINGS['MODE_DEBUG']) parser.add_argument('--config', action = 'store', dest = 'JSON_CONFIG', help = 'Configure from a JSON file rather than the command line', default = '') + parser.add_argument('--static-config', action = 'store', dest = 'STATIC_CONFIG', help = 'Configure leases from a json file rather than the command line', default = '') parser.add_argument('--syslog', action = 'store', dest = 'SYSLOG_SERVER', help = 'Syslog server', default = SETTINGS['SYSLOG_SERVER']) parser.add_argument('--syslog-port', action = 'store', dest = 'SYSLOG_PORT', help = 'Syslog server port', default = SETTINGS['SYSLOG_PORT']) + # DHCP server arguments - exclusive = parser.add_mutually_exclusive_group(required = False) + dhcp_group = parser.add_argument_group(title = 'DHCP', description = 'Arguments relevant to the DHCP server') + exclusive = dhcp_group.add_mutually_exclusive_group(required = False) + exclusive.add_argument('--no-dhcp', action = 'store_false', dest = 'USE_DHCP', help = 'Disable built-in DHCP server', default = not SETTINGS['USE_DHCP']) exclusive.add_argument('--dhcp', action = 'store_true', dest = 'USE_DHCP', help = 'Enable built-in DHCP server', default = SETTINGS['USE_DHCP']) exclusive.add_argument('--dhcp-proxy', action = 'store_true', dest = 'DHCP_MODE_PROXY', help = 'Enable built-in DHCP server in proxy mode (implies --dhcp)', default = SETTINGS['DHCP_MODE_PROXY']) - parser.add_argument('-s', '--dhcp-server-ip', action = 'store', dest = 'DHCP_SERVER_IP', help = 'DHCP Server IP', default = SETTINGS['DHCP_SERVER_IP']) - parser.add_argument('-p', '--dhcp-server-port', action = 'store', dest = 'DHCP_SERVER_PORT', help = 'DHCP Server Port', default = SETTINGS['DHCP_SERVER_PORT']) - parser.add_argument('-b', '--dhcp-begin', action = 'store', dest = 'DHCP_OFFER_BEGIN', help = 'DHCP lease range start', default = SETTINGS['DHCP_OFFER_BEGIN']) - parser.add_argument('-e', '--dhcp-end', action = 'store', dest = 'DHCP_OFFER_END', help = 'DHCP lease range end', default = SETTINGS['DHCP_OFFER_END']) - parser.add_argument('-n', '--dhcp-subnet', action = 'store', dest = 'DHCP_SUBNET', help = 'DHCP lease subnet', default = SETTINGS['DHCP_SUBNET']) - parser.add_argument('-r', '--dhcp-router', action = 'store', dest = 'DHCP_ROUTER', help = 'DHCP lease router', default = SETTINGS['DHCP_ROUTER']) - parser.add_argument('-d', '--dhcp-dns', action = 'store', dest = 'DHCP_DNS', help = 'DHCP lease DNS server', default = SETTINGS['DHCP_DNS']) - parser.add_argument('-c', '--dhcp-broadcast', action = 'store', dest = 'DHCP_BROADCAST', help = 'DHCP broadcast address', default = SETTINGS['DHCP_BROADCAST']) - parser.add_argument('-f', '--dhcp-fileserver', action = 'store', dest = 'DHCP_FILESERVER', help = 'DHCP fileserver IP', default = SETTINGS['DHCP_FILESERVER']) + + dhcp_group.add_argument('--dhcp-server-ip', action = 'store', dest = 'DHCP_SERVER_IP', help = 'DHCP Server IP', default = SETTINGS['DHCP_SERVER_IP']) + dhcp_group.add_argument('--dhcp-server-port', action = 'store', dest = 'DHCP_SERVER_PORT', help = 'DHCP Server Port', default = SETTINGS['DHCP_SERVER_PORT']) + dhcp_group.add_argument('--dhcp-begin', action = 'store', dest = 'DHCP_OFFER_BEGIN', help = 'DHCP lease range start', default = SETTINGS['DHCP_OFFER_BEGIN']) + dhcp_group.add_argument('--dhcp-end', action = 'store', dest = 'DHCP_OFFER_END', help = 'DHCP lease range end', default = SETTINGS['DHCP_OFFER_END']) + dhcp_group.add_argument('--dhcp-subnet', action = 'store', dest = 'DHCP_SUBNET', help = 'DHCP lease subnet', default = SETTINGS['DHCP_SUBNET']) + dhcp_group.add_argument('--dhcp-router', action = 'store', dest = 'DHCP_ROUTER', help = 'DHCP lease router', default = SETTINGS['DHCP_ROUTER']) + dhcp_group.add_argument('--dhcp-dns', action = 'store', dest = 'DHCP_DNS', help = 'DHCP lease DNS server', default = SETTINGS['DHCP_DNS']) + dhcp_group.add_argument('--dhcp-broadcast', action = 'store', dest = 'DHCP_BROADCAST', help = 'DHCP broadcast address', default = SETTINGS['DHCP_BROADCAST']) + dhcp_group.add_argument('--dhcp-fileserver', action = 'store', dest = 'DHCP_FILESERVER', help = 'DHCP fileserver IP', default = SETTINGS['DHCP_FILESERVER']) + dhcp_group.add_argument('--dhcp-whitelist', action = 'store_true', dest = 'DHCP_WHITELIST', help = 'Only respond to DHCP clients present in --static-config', default = False) # network boot directory and file name arguments - parser.add_argument('-a', '--netboot-dir', action = 'store', dest = 'NETBOOT_DIR', help = 'Local file serve directory', default = SETTINGS['NETBOOT_DIR']) - parser.add_argument('-i', '--netboot-file', action = 'store', dest = 'NETBOOT_FILE', help = 'PXE boot file name (after iPXE if --ipxe)', default = SETTINGS['NETBOOT_FILE']) + parser.add_argument('--netboot-dir', action = 'store', dest = 'NETBOOT_DIR', help = 'Local file serve directory', default = SETTINGS['NETBOOT_DIR']) + parser.add_argument('--netboot-file', action = 'store', dest = 'NETBOOT_FILE', help = 'PXE boot file name (after iPXE if --ipxe)', default = SETTINGS['NETBOOT_FILE']) + + # NBD server arguments + nbd_group = parser.add_argument_group(title = 'Network Block Device', description = 'Arguments relevant to the NBD server') + nbd_group.add_argument('--nbd', action = 'store', dest = 'NBD_BLOCK_DEVICE', help = 'Enable the NDB server with a specific block device (Can be a disk image)', default = SETTINGS['NBD_BLOCK_DEVICE']) + nbd_group.add_argument('--nbd-write', action = 'store_true', dest = 'NBD_WRITE', help = 'Enable writes on the NBD device', default = SETTINGS['NBD_WRITE']) + nbd_group.add_argument('--nbd-cow', action = 'store_true', dest = 'NBD_COW', help = 'Enable copy-on-write for the NBD device (Non-persistent changes)', default = SETTINGS['NBD_COW']) + nbd_group.add_argument('--nbd-cow-in-mem', action = 'store_true', dest = 'NBD_COW_IN_MEM', help = 'Store copy-on-write pages in memory', default = SETTINGS['NBD_COW_IN_MEM']) + nbd_group.add_argument('--nbd-copy-to-ram', action = 'store_true', dest = 'NBD_COPY_TO_RAM', help = 'Copy the NBD device to memory before serving clients', default = SETTINGS['NBD_COPY_TO_RAM']) + nbd_group.add_argument('--nbd-server', action = 'store', dest = 'NBD_SERVER_IP', help = 'NBD Server IP', default = SETTINGS['NBD_SERVER_IP']) + nbd_group.add_argument('--nbd-port', action = 'store', dest = 'NBD_PORT', help = 'NBD Server Port', default = SETTINGS['NBD_PORT']) return parser.parse_args() +def do_debug(service): + return ((service in args.MODE_DEBUG.lower() + or 'all' in args.MODE_DEBUG.lower()) + and '-{0}'.format(service) not in args.MODE_DEBUG.lower()) + if __name__ == '__main__': try: # warn the user that they are starting PyPXE as non-root user @@ -92,6 +131,21 @@ def parse_cli_arguments(): SETTINGS.update(loaded_config) # update settings with JSON config args = parse_cli_arguments() # re-parse, CLI options take precedence + # ideally this would be in dhcp itself, but the chroot below *probably* + # breaks the ability to open the config file. + if args.STATIC_CONFIG: + try: + static_config = open(args.STATIC_CONFIG, 'rb') + except IOError: + sys.exit("Failed to open {0}".format(args.STATIC_CONFIG)) + try: + loaded_statics = json.load(static_config) + static_config.close() + except ValueError: + sys.exit("{0} does not contain valid json".format(args.STATIC_CONFIG)) + else: + loaded_statics = dict() + # setup main logger sys_logger = logging.getLogger('PyPXE') if args.SYSLOG_SERVER: @@ -121,7 +175,13 @@ def parse_cli_arguments(): else: args.NETBOOT_FILE = 'boot.http.ipxe' - # serve all files from one directory + if args.NBD_WRITE and not args.NBD_COW: + sys_logger.warning('NBD Write enabled but copy-on-write is not. Multiple clients may cause corruption') + + if args.NBD_COW_IN_MEM or args.NBD_COPY_TO_RAM: + sys_logger.warning('NBD cowinmem and copytoram can cause high RAM usage') + + #serve all files from one directory os.chdir (args.NETBOOT_DIR) # make a list of running threads for each service @@ -135,7 +195,7 @@ def parse_cli_arguments(): sys_logger.info('Starting TFTP server...') # setup the thread - tftp_server = tftp.TFTPD(mode_debug = ('tftp' in args.MODE_DEBUG.lower() or 'all' in args.MODE_DEBUG.lower()), logger = tftp_logger) + tftp_server = tftp.TFTPD(mode_debug = do_debug('tftp'), logger = tftp_logger) tftpd = threading.Thread(target = tftp_server.listen) tftpd.daemon = True tftpd.start() @@ -153,21 +213,23 @@ def parse_cli_arguments(): # setup the thread dhcp_server = dhcp.DHCPD( - ip = args.DHCP_SERVER_IP, - port = args.DHCP_SERVER_PORT, - offer_from = args.DHCP_OFFER_BEGIN, - offer_to = args.DHCP_OFFER_END, - subnet_mask = args.DHCP_SUBNET, - router = args.DHCP_ROUTER, - dns_server = args.DHCP_DNS, - broadcast = args.DHCP_BROADCAST, - file_server = args.DHCP_FILESERVER, - file_name = args.NETBOOT_FILE, - use_ipxe = args.USE_IPXE, - use_http = args.USE_HTTP, - mode_proxy = args.DHCP_MODE_PROXY, - mode_debug = ('dhcp' in args.MODE_DEBUG.lower() or 'all' in args.MODE_DEBUG.lower()), - logger = dhcp_logger) + ip = args.DHCP_SERVER_IP, + port = args.DHCP_SERVER_PORT, + offer_from = args.DHCP_OFFER_BEGIN, + offer_to = args.DHCP_OFFER_END, + subnet_mask = args.DHCP_SUBNET, + router = args.DHCP_ROUTER, + dns_server = args.DHCP_DNS, + broadcast = args.DHCP_BROADCAST, + file_server = args.DHCP_FILESERVER, + file_name = args.NETBOOT_FILE, + use_ipxe = args.USE_IPXE, + use_http = args.USE_HTTP, + mode_proxy = args.DHCP_MODE_PROXY, + mode_debug = do_debug('dhcp'), + whitelist = args.DHCP_WHITELIST, + static_config = loaded_statics, + logger = dhcp_logger) dhcpd = threading.Thread(target = dhcp_server.listen) dhcpd.daemon = True dhcpd.start() @@ -181,12 +243,32 @@ def parse_cli_arguments(): sys_logger.info('Starting HTTP server...') # setup the thread - http_server = http.HTTPD(mode_debug = ('http' in args.MODE_DEBUG.lower() or 'all' in args.MODE_DEBUG.lower()), logger = http_logger) + http_server = http.HTTPD(mode_debug = do_debug('http'), logger = http_logger) httpd = threading.Thread(target = http_server.listen) httpd.daemon = True httpd.start() running_services.append(httpd) + # configure/start NBD server + if args.NBD_BLOCK_DEVICE: + # setup NBD logger + nbd_logger = sys_logger.getChild('NBD') + sys_logger.info('Starting NBD server...') + nbd_server = nbd.NBD( + block_device = args.NBD_BLOCK_DEVICE, + write = args.NBD_WRITE, + cow = args.NBD_COW, + in_mem = args.NBD_COW_IN_MEM, + copy_to_ram = args.NBD_COPY_TO_RAM, + ip = args.NBD_SERVER_IP, + port = args.NBD_PORT, + mode_debug = ('nbd' in args.MODE_DEBUG.lower() or 'all' in args.MODE_DEBUG.lower()), + logger = nbd_logger) + nbd = threading.Thread(target = nbd_server.listen) + nbd.daemon = True + nbd.start() + running_services.append(nbd) + sys_logger.info('PyPXE successfully initialized and running!') while map(lambda x: x.isAlive(), running_services): diff --git a/pypxe/__init__.py b/pypxe/__init__.py index fcb6b5d..6d5e09d 100644 --- a/pypxe/__init__.py +++ b/pypxe/__init__.py @@ -1 +1 @@ -__version__ = '1.5' +__version__ = '1.6' diff --git a/pypxe/dhcp.py b/pypxe/dhcp.py index 89156bc..46441de 100644 --- a/pypxe/dhcp.py +++ b/pypxe/dhcp.py @@ -14,6 +14,7 @@ class OutOfLeasesError(Exception): pass + class DHCPD: ''' This class implements a DHCP Server, limited to PXE options. @@ -41,9 +42,11 @@ def __init__(self, **server_settings): self.ipxe = server_settings.get('use_ipxe', False) self.http = server_settings.get('use_http', False) self.mode_proxy = server_settings.get('mode_proxy', False) # ProxyDHCP mode + self.static_config = server_settings.get('static_config', dict()) + self.whitelist = server_settings.get('whitelist', False) self.mode_debug = server_settings.get('mode_debug', False) # debug mode - self.magic = struct.pack('!I', 0x63825363) # magic cookie self.logger = server_settings.get('logger', None) + self.magic = struct.pack('!I', 0x63825363) # magic cookie # setup logger if self.logger == None: @@ -68,14 +71,18 @@ def __init__(self, **server_settings): # debug info for ProxyDHCP mode if not self.mode_proxy: - self.logger.debug('DHCP Lease Range: {0} - {1}'.format(self.offer_from, self.offer_to)) - self.logger.debug('DHCP Subnet Mask: {0}'.format(self.subnet_mask)) - self.logger.debug('DHCP Router: {0}'.format(self.router)) - self.logger.debug('DHCP DNS Server: {0}'.format(self.dns_server)) - self.logger.debug('DHCP Broadcast Address: {0}'.format(self.broadcast)) - - self.logger.debug('DHCP File Server IP: {0}'.format(self.file_server)) - self.logger.debug('DHCP File Name: {0}'.format(self.file_name)) + self.logger.debug('Lease Range: {0} - {1}'.format(self.offer_from, self.offer_to)) + self.logger.debug('Subnet Mask: {0}'.format(self.subnet_mask)) + self.logger.debug('Router: {0}'.format(self.router)) + self.logger.debug('DNS Server: {0}'.format(self.dns_server)) + self.logger.debug('Broadcast Address: {0}'.format(self.broadcast)) + + if self.static_config: + self.logger.debug('Using Static Leasing') + self.logger.debug('Using Static Leasing Whitelist: {0}'.format(self.whitelist)) + + self.logger.debug('File Server IP: {0}'.format(self.file_server)) + self.logger.debug('File Name: {0}'.format(self.file_name)) self.logger.debug('ProxyDHCP Mode: {0}'.format(self.mode_proxy)) self.logger.debug('Using iPXE: {0}'.format(self.ipxe)) self.logger.debug('Using HTTP Server: {0}'.format(self.http)) @@ -88,6 +95,12 @@ def __init__(self, **server_settings): # key is MAC self.leases = defaultdict(lambda: {'ip': '', 'expire': 0, 'ipxe': self.ipxe}) + def get_namespaced_static(self, path, fallback = {}): + statics = self.static_config + for child in path.split('.'): + statics = statics.get(child, {}) + return statics if statics else fallback + def next_ip(self): ''' This method returns the next unleased IP from range; @@ -142,7 +155,7 @@ def tlv_parse(self, raw): ret[tag] = [value] return ret - def print_mac(self, mac): + def get_mac(self, mac): ''' This method converts the MAC Address from binary to human-readable format for logging. @@ -164,10 +177,11 @@ def craft_header(self, message): if self.leases[client_mac]['ip']: # OFFER offer = self.leases[client_mac]['ip'] else: # ACK - offer = self.next_ip() + offer = self.get_namespaced_static('dhcp.binding.{0}.ipaddr'.format(self.get_mac(client_mac))) + offer = offer if offer else self.next_ip() self.leases[client_mac]['ip'] = offer self.leases[client_mac]['expire'] = time() + 86400 - self.logger.debug('New Assignment - MAC: {0} -> IP: {1}'.format(self.print_mac(client_mac), self.leases[client_mac]['ip'])) + self.logger.debug('New Assignment - MAC: {0} -> IP: {1}'.format(self.get_mac(client_mac), self.leases[client_mac]['ip'])) response += socket.inet_aton(offer) # yiaddr else: response += socket.inet_aton('0.0.0.0') @@ -191,14 +205,18 @@ def craft_options(self, opt53, client_mac): opt53: 2 - DHCPOFFER 5 - DHCPACK - (See RFC2132 9.6) + See RFC2132 9.6 for details. ''' response = self.tlv_encode(53, chr(opt53)) # message type, OFFER response += self.tlv_encode(54, socket.inet_aton(self.ip)) # DHCP Server if not self.mode_proxy: - response += self.tlv_encode(1, socket.inet_aton(self.subnet_mask)) # Subnet Mask - response += self.tlv_encode(3, socket.inet_aton(self.router)) # Router - response += self.tlv_encode(6, socket.inet_aton(self.dns_server)) # DNS + subnet_mask = self.get_namespaced_static('dhcp.binding.{0}.subnet'.format(self.get_mac(client_mac)), self.subnet_mask) + response += self.tlv_encode(1, socket.inet_aton(subnet_mask)) # subnet mask + router = self.get_namespaced_static('dhcp.binding.{0}.router'.format(self.get_mac(client_mac)), self.router) + response += self.tlv_encode(3, socket.inet_aton(router)) # router + dns_server = self.get_namespaced_static('dhcp.binding.{0}.dns'.format(self.get_mac(client_mac)), [self.dns_server]) + dns_server = ''.join([socket.inet_aton(i) for i in dns_server]) + response += self.tlv_encode(6, dns_server) response += self.tlv_encode(51, struct.pack('!I', 86400)) # lease time # TFTP Server OR HTTP Server; if iPXE, need both @@ -207,16 +225,16 @@ def craft_options(self, opt53, client_mac): # file_name null terminated if not self.ipxe or not self.leases[client_mac]['ipxe']: # http://www.syslinux.org/wiki/index.php/PXELINUX#UEFI - if 93 in self.options and not self.force_file_name: - [arch] = struct.unpack("!H", self.options[93][0]) + if 93 in self.leases[client_mac]['options'] and not self.force_file_name: + [arch] = struct.unpack("!H", self.leases[client_mac]['options'][93][0]) if arch == 0: # BIOS/default - response += self.tlv_encode(67, "pxelinux.0" + chr(0)) + response += self.tlv_encode(67, 'pxelinux.0' + chr(0)) elif arch == 6: # EFI IA32 - response += self.tlv_encode(67, "syslinux.efi32" + chr(0)) - elif arch == 7: # EFI BC, x86-64 according to link above - response += self.tlv_encode(67, "syslinux.efi64" + chr(0)) + response += self.tlv_encode(67, 'syslinux.efi32' + chr(0)) + elif arch == 7: # EFI BC, x86-64 (according to the above link) + response += self.tlv_encode(67, 'syslinux.efi64' + chr(0)) elif arch == 9: # EFI x86-64 - response += self.tlv_encode(67, "syslinux.efi64" + chr(0)) + response += self.tlv_encode(67, 'syslinux.efi64' + chr(0)) else: response += self.tlv_encode(67, self.file_name + chr(0)) else: @@ -263,9 +281,12 @@ def dhcp_ack(self, message): self.logger.debug('<--END RESPONSE-->') self.sock.sendto(response, (self.broadcast, 68)) - def validate_req(self): + def validate_req(self, client_mac): # client request is valid only if contains Vendor-Class = PXEClient - if 60 in self.options and 'PXEClient' in self.options[60][0]: + if self.whitelist and self.get_mac(client_mac) not in self.get_namespaced_static('dhcp.binding'): + self.logger.debug('Non-whitelisted client request received') + return False + if 60 in self.leases[client_mac]['options'] and 'PXEClient' in self.leases[client_mac]['options'][60][0]: self.logger.debug('PXE client request received') return True if self.mode_debug: @@ -276,25 +297,25 @@ def listen(self): '''Main listen loop.''' while True: message, address = self.sock.recvfrom(1024) - client_mac = struct.unpack('!28x6s', message[:34]) + [client_mac] = struct.unpack('!28x6s', message[:34]) self.logger.debug('Received message') self.logger.debug('<--BEGIN MESSAGE-->') self.logger.debug('{0}'.format(repr(message))) self.logger.debug('<--END MESSAGE-->') - self.options = self.tlv_parse(message[240:]) + self.leases[client_mac]['options'] = self.tlv_parse(message[240:]) self.logger.debug('Parsed received options') self.logger.debug('<--BEGIN OPTIONS-->') - self.logger.debug('{0}'.format(repr(self.options))) + self.logger.debug('{0}'.format(repr(self.leases[client_mac]['options']))) self.logger.debug('<--END OPTIONS-->') - if not self.validate_req(): + if not self.validate_req(client_mac): continue - type = ord(self.options[53][0]) # see RFC2131, page 10 + type = ord(self.leases[client_mac]['options'][53][0]) # see RFC2131, page 10 if type == 1: self.logger.debug('Received DHCPOFFER') try: self.dhcp_offer(message) except OutOfLeasesError: - self.logger.critical("Ran out of leases") + self.logger.critical('Ran out of leases') elif type == 3 and address[0] == '0.0.0.0' and not self.mode_proxy: self.logger.debug('Received DHCPACK') self.dhcp_ack(message) diff --git a/pypxe/http.py b/pypxe/http.py index d1f97fc..2ef325a 100644 --- a/pypxe/http.py +++ b/pypxe/http.py @@ -45,9 +45,9 @@ def __init__(self, **server_settings): os.chroot ('.') self.logger.debug('NOTICE: HTTP server started in debug mode. HTTP server is using the following:') - self.logger.debug('HTTP Server IP: {0}'.format(self.ip)) - self.logger.debug('HTTP Server Port: {0}'.format(self.port)) - self.logger.debug('HTTP Network Boot Directory: {0}'.format(self.netboot_directory)) + self.logger.debug('Server IP: {0}'.format(self.ip)) + self.logger.debug('Server Port: {0}'.format(self.port)) + self.logger.debug('Network Boot Directory: {0}'.format(self.netboot_directory)) def handle_request(self, connection, addr): '''This method handles HTTP request.''' diff --git a/pypxe/nbd/__init__.py b/pypxe/nbd/__init__.py new file mode 100644 index 0000000..2c0c943 --- /dev/null +++ b/pypxe/nbd/__init__.py @@ -0,0 +1 @@ +from nbd import * diff --git a/pypxe/nbd/nbd.py b/pypxe/nbd/nbd.py new file mode 100644 index 0000000..8a1ad53 --- /dev/null +++ b/pypxe/nbd/nbd.py @@ -0,0 +1,167 @@ +import logging +import socket +import struct +import threading +import sys +import os +import io +import writes + +class NBD: + def __init__(self, **server_settings): + self.bd = server_settings.get('block_device', '') + self.write = server_settings.get('write', False) # write? + self.cow = server_settings.get('cow', True) # COW is the safe default + self.in_mem = server_settings.get('in_mem', False) + self.copy_to_ram = server_settings.get('copy_to_ram', False) + self.ip = server_settings.get('ip', '0.0.0.0') + self.port = server_settings.get('port', 10809) + self.mode_debug = server_settings.get('mode_debug', False) # debug mode + self.logger = server_settings.get('logger', None) + + # setup logger + if self.logger == None: + self.logger = logging.getLogger('NBD') + handler = logging.StreamHandler() + formatter = logging.Formatter('%(asctime)s %(name)s [%(levelname)s] %(message)s') + handler.setFormatter(formatter) + self.logger.addHandler(handler) + + if self.mode_debug: + self.logger.setLevel(logging.DEBUG) + + self.logger.debug('NOTICE: NBD server started in debug mode. NBD server is using the following:') + self.logger.debug('Server IP: {0}'.format(self.ip)) + self.logger.debug('Server Port: {0}'.format(self.port)) + self.logger.debug('Block Device: {0}'.format(self.bd)) + self.logger.debug('Block Device Writes: {0}'.format(self.write)) + self.logger.debug('Block Write Method: {0} ({1})'.format("Copy-On-Write" if self.cow else 'File', 'Memory' if self.in_mem else 'Disk')) + + if self.mode_debug: + self.logger.setLevel(logging.DEBUG) + + self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind((self.ip, self.port)) + self.sock.listen(4) + + # if we have COW on, we write elsewhere so we don't need write ability + self.openbd = open(self.bd, 'r+b' if self.write and not self.cow else 'rb') + # go to EOF + self.openbd.seek(0, 2) + # we need this when clients mount us + self.bdsize = self.openbd.tell() + # go back to start + self.openbd.seek(0) + if self.copy_to_ram and self.cow: + self.logger.info('Starting copying {0} to RAM'.format(self.bd)) + self.openbd = io.BytesIO(self.openbd.read()) + self.logger.info('Finished copying {0} to RAM'.format(self.bd)) + + def send_reply(self, conn, addr, code, data): + '''Send a reply with magic, only used for error codes.''' + reply = struct.pack('!Q', 0x3e889045565a9) + reply += struct.pack('!I', code) + reply += struct.pack('!I', len(data)) + reply += data + conn.send(reply) + + def handshake(self, conn, addr): + '''Initiate the connection, server sends first.''' + # mostly taken from https://github.com/yoe/nbd/blob/master/nbd-server.c + conn.send('NBDMAGIC') + # 0x49484156454F5054 + conn.send('IHAVEOPT') + # NBD_FLAG_FIXED_NEWSTYLE | NBD_FLAG_NO_ZEROES + conn.send(struct.pack('!H', 3)) + + [cflags] = struct.unpack('!I', conn.recv(4)) + + op = 0 + while op != 1: # NBD_OPT_EXPORT_NAME + [magic] = struct.unpack("!Q", conn.recv(8)) + [op] = struct.unpack("!I", conn.recv(4)) + if op != 1: + # NBD_REP_ERR_UNSUP + self.send_reply(conn, addr, 2 ** 31 + 1, '') + + [namelen] = struct.unpack('!I', conn.recv(4)) + name = conn.recv(namelen) + if name != self.bd: + conn.close() + return 1 + + self.logger.debug('Received request for {0} from {1}'.format(name, addr)) + + # size of export + exportinfo = struct.pack('!Q', self.bdsize) + flags = (0 if self.write else 2) # readonly? + exportinfo += struct.pack('!H', flags) + exportinfo += chr(0) * (0 if (cflags & 2) else 124) + conn.send(exportinfo) + + def handle_client(self, conn, addr, seeklock): + '''Handle all client actions, R/W/Disconnect''' + ret = self.handshake(conn, addr) + if ret: return # client did something wrong, so we closed them + + FS = writes.write(self.cow, self.in_mem)(addr, self.openbd, self.logger, seeklock) + + while True: + conn.recv(4) + try: + [opcode, handle, offset, length] = struct.unpack('!IQQI', conn.recv(24, socket.MSG_WAITALL)) + except struct.error: + # client sent us something malformed, or gave up (module not loaded) + continue + if opcode not in (0, 1, 2): + # NBD_REP_ERR_UNSUP + self.send_reply(conn, addr, 2 ** 31 + 1, '') + continue + if opcode == 0: # READ + data = FS.read(offset, length) + + response = struct.pack('!I', 0x67446698) + response += struct.pack('!I', 0) # error + response += struct.pack('!Q', handle) + conn.send(response) + conn.send(data) + + elif opcode == 1: # WRITE + # WAIT because if there's any lag at all we don't get the whole + # thing, we don't write the whole thing, and then we break + # trying to parse the rest of the data + data = conn.recv(length, socket.MSG_WAITALL) + FS.write(offset, data) + + response = struct.pack('!I', 0x67446698) + response += struct.pack('!I', 0) # error + response += struct.pack('!Q', handle) + conn.send(response) + + elif opcode == 2: # DISCONNECT + # delete COW diff + conn.close() + self.logger.debug('{0} disconnected'.format(addr)) + return + + def listen(self): + '''This method is the main loop that listens for requests.''' + seeklock = threading.Lock() + cowfiles = [] + while True: + try: + conn, addr = self.sock.accept() + # split off on a thread, allows us to handle multiple clients + dispatch = threading.Thread(target = self.handle_client, args = (conn, addr, seeklock)) + # clients don't necessarily close the TCP connection + # so we use this to kill the program on ctrl-c + dispatch.daemon = True + dispatch.start() + # this is for the cleanup at the end. Will need clarifying + # if MemCOW + if self.cow and not self.in_mem: + cowfiles.append('PyPXE_NBD_COW_{addr[0]}_{addr[1]}'.format(addr = addr)) + except KeyboardInterrupt: + map(os.remove, cowfiles) + return diff --git a/pypxe/nbd/writes.py b/pypxe/nbd/writes.py new file mode 100644 index 0000000..30cd787 --- /dev/null +++ b/pypxe/nbd/writes.py @@ -0,0 +1,140 @@ +import io + +class COW: + def basepages(self, offset, length): + # basepages is (page base addr, inter page offset, length of data in page) + # it's unlikely we'll ever need sub 4096 reads, but just in case + basepages = [] + + # first chunk, not necessarily at page boundary + basepages.append((offset - (offset % 4096), offset % 4096, 4096 - (offset % 4096))) + length -= 4096 - (offset % 4096) + offset += 4096 + + # all following FULL chunks, definate page boundary and full size + while length >= 4096: + basepages.append((offset, 0, 4096)) + length -= 4096 + offset += 4096 + + # final non-full chunk, definate offset, variable length + if length > 0: + basepages.append((offset, 0, length)) + + return basepages + + def read(self, offset, length): + basepages = self.basepages(offset, length) + + # this probably wants to be 2nd debug level + self.logger.debug('{0} reading {1} bytes from {2}. Pages: {3}'.format(self.addr, length, hex(offset), len(basepages))) + + data = '' + for major, minor, length in basepages: + if major in self.pages: + # major is the nth page in the file + off = self.pages.index(major) + self.fh.seek(off * 4096 + minor) + data += self.fh.read(length) + else: + # This is a race condition. If another thread seeks after we've + # seeked, but before we read, the seek is changed and the data + # is wrong. + # Lock is shared between all clients. Only applies to imagefd, + # self.fd is unique and per client. + self.seek_lock.acquire() + self.imagefd.seek(major + minor) + data += self.imagefd.read(length) + self.seek_lock.release() + return data + + def write(self, offset, data): + basepages = self.basepages(offset, len(data)) + + # this probably wants to be 2nd debug level + self.logger.debug('{0} writing {1} bytes to {2}. Pages: {3}'.format(self.addr, len(data), hex(offset), len(basepages))) + + for major, minor, length in basepages: + if major in self.pages: + # we already have a copied page, so we can just overwrite it + self.fh.seek(major + minor) + self.fh.write(data[:length]) + data = data[length:] + else: + # we don't have this page, so copy it first. then add to the list + self.seek_lock.acquire() + # on the page boundary + self.imagefd.seek(major) + cpdata = self.imagefd.read(4096) + self.seek_lock.release() + # append to EOF + self.fh.seek(0, 2) + self.fh.write(cpdata) + self.pages.append(major) + # we've got a copy of the page now, we just need to write it + off = self.pages.index(major) + self.fh.seek(off * 4096 + minor) + self.fh.write(data[:length]) + data = data[length:] + +class DiskCOW(COW): + def __init__(self, addr, imagefd, logger, seek_lock): + # optional argset for disk diff path + self.addr = addr + self.imagefd = imagefd + self.seek_lock = seek_lock + self.logger = logger.getChild('FS') + self.logger.debug('Copy-On-Write for {addr} in PyPXE_NBD_COW_{addr[0]}_{addr[1]}'.format(addr = addr)) + + # never want readonly cow, also definately creating file + self.fh = open('PyPXE_NBD_COW_{addr[0]}_{addr[1]}'.format(addr = addr), 'w+b') + # pages is a list of the addresses for which we have different pages + # should all be multiples of 4096 + self.pages = [] + +class MemCOW(COW): + def __init__(self, addr, imagefd, logger, seek_lock): + self.addr = addr + self.imagefd = imagefd + self.seek_lock = seek_lock + self.logger = logger.getChild('FS') + self.logger.debug('Copy-On-Write for {0} in Memory'.format(addr)) + + # BytesIO looks exactly the same as a file, perfect for in memory disk + self.fh = io.BytesIO() + # pages is a list of the addresses for which we have different pages + # should all be multiples of 4096 + self.pages = [] + +class RW: + def __init__(self, addr, imagefd, logger, seek_lock): + self.addr = addr + self.seek_lock = seek_lock + self.imagefd = imagefd + self.logger = logger.getChild('FS') + self.logger.debug('File for {0}'.format(addr)) + + def read(self, offset, length): + self.logger.debug('{0} reading {1} bytes from [{2}]'.format(self.addr, length, hex(offset))) + # see COW.read() for lock reason + self.seek_lock.acquire() + self.imagefd.seek(offset) + data = self.imagefd.read(length) + self.seek_lock.release() + return data + + def write(self, offset, data): + self.logger.debug('{0} writing {1} bytes to {2}'.format(self.addr, len(data), hex(offset))) + self.seek_lock.acquire() + self.imagefd.seek(offset) + self.imagefd.write(data) + self.seek_lock.release() + +def write(cow, in_mem): + '''Class signatures are identical so we can transparently use either.''' + if cow and in_mem: + return MemCOW + elif cow and not in_mem: + return DiskCOW + else: + return RW diff --git a/pypxe/tftp.py b/pypxe/tftp.py index ac4c1f4..d436da5 100644 --- a/pypxe/tftp.py +++ b/pypxe/tftp.py @@ -98,7 +98,7 @@ def parse_options(self): self.blksize = options.get('blksize', self.blksize) self.lastblock = math.ceil(self.filesize / float(self.blksize)) self.tsize = True if 'tsize' in options else False - if self.filesize > (2**16)*self.blksize: + if self.filesize > (2 ** 16) * self.blksize: self.logger.warning('Request too big, attempting transfer anyway.') self.logger.debug('Details: Filesize {0} is too big for blksize {1}.'.format(self.filesize, self.blksize)) @@ -148,7 +148,7 @@ def newRequest(self): self.send_block() return - # we got some options, so ack those first + # we got some options so ACK those first self.reply_options() def sendError(self, code = 1, message = 'File Not Found', filename = ''): @@ -182,7 +182,7 @@ def complete(self): try: self.fh.close() except AttributeError: - # we have not opened yet, or file-not-found + # we have not opened yet or file-not-found pass self.sock.close() self.dead = True @@ -238,9 +238,9 @@ def __init__(self, **server_settings): self.logger.setLevel(logging.DEBUG) self.logger.debug('NOTICE: TFTP server started in debug mode. TFTP server is using the following:') - self.logger.debug('TFTP Server IP: {0}'.format(self.ip)) - self.logger.debug('TFTP Server Port: {0}'.format(self.port)) - self.logger.debug('TFTP Network Boot Directory: {0}'.format(self.netbook_directory)) + self.logger.debug('Server IP: {0}'.format(self.ip)) + self.logger.debug('Server Port: {0}'.format(self.port)) + self.logger.debug('Network Boot Directory: {0}'.format(self.netbook_directory)) self.ongoing = []